Use fetch_url for image data

3 files changed, 44 insertions, 16 deletions

diff --git a/gsxws/parts.py b/gsxws/parts.py
index 6f9fb02..0f3691b 100644
--- a/gsxws/parts.py
+++ b/gsxws/parts.py
@@ -1,10 +1,8 @@
 # -*- coding: utf-8 -*-
 
-import urllib
-import tempfile
-
-from lookups import Lookup
-from core import GsxObject, GsxError
+from .lookups import Lookup
+from .utils import fetch_url
+from .core import GsxObject, GsxError
 
 REASON_CODES = (
     ('A', 'Part not needed'),
@@ -35,11 +33,9 @@ class Part(GsxObject):
             raise GsxError("Cannot fetch part image without part number")
 
         image = "%s_350_350.gif" % self.partNumber
-        url = IMAGE_URL % image
-        tmpfile = tempfile.mkstemp(suffix=image)
 
         try:
-            return urllib.urlretrieve(url, tmpfile[1])[0]
+            return fetch_url(IMAGE_URL % image)
         except Exception as e:
             raise GsxError("Failed to fetch part image: %s" % e)
 
@@ -48,7 +44,7 @@ if __name__ == '__main__':
     import sys
     import doctest
     import logging
-    from core import connect
+    from .core import connect
     logging.basicConfig(level=logging.DEBUG)
     connect(*sys.argv[1:])
     doctest.testmod()
diff --git a/gsxws/products.py b/gsxws/products.py
index 2c2745a..a656f98 100644
--- a/gsxws/products.py
+++ b/gsxws/products.py
@@ -4,11 +4,11 @@ https://gsxwsut.apple.com/apidocs/ut/html/WSAPIChangeLog.html?user=asp
 """
 
 import re
-import urllib
 
-from lookups import Lookup
-from diagnostics import Diagnostics
-from core import GsxObject, GsxError, validate
+from .utils import fetch_url
+from .lookups import Lookup
+from .diagnostics import Diagnostics
+from .core import GsxObject, GsxError, validate
 
 
 def models():
@@ -19,7 +19,8 @@ def models():
     import os
     import yaml
     filepath = os.path.join(os.path.dirname(__file__), "products.yaml")
-    return yaml.load(open(filepath, 'r'))
+    with open(filepath, 'r') as f:
+        return yaml.load(f)
 
 
 class Product(object):
@@ -151,7 +152,7 @@ class Product(object):
             raise GsxError("No URL to fetch product image")
 
         try:
-            return urllib.urlretrieve(url)[0]
+            return fetch_url(url)
         except Exception as e:
             raise GsxError("Failed to fetch product image: %s" % e)
 
@@ -262,7 +263,7 @@ if __name__ == '__main__':
     import sys
     import doctest
     import logging
-    from core import connect
+    from .core import connect
     logging.basicConfig(level=logging.DEBUG)
     connect(*sys.argv[1:])
     doctest.testmod()
diff --git a/gsxws/utils.py b/gsxws/utils.py
new file mode 100644
index 0000000..9d71e83
--- /dev/null
+++ b/gsxws/utils.py
@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+
+import re
+import tempfile
+import requests
+
+
+def fetch_url(url):
+    """Downloads acceptable URL to temporary file
+    """
+    ext = None
+    ALLOWED = ['jpg', 'jpeg', 'png', 'pdf', 'gif',]
+
+    try:
+        ext = re.search(r'\.([a-zA-Z]{3,})$', url).group(1)
+    except AttributeError:
+        raise ValueError('Cannot determine file extension of URL %s' % url)
+
+    ext = ext.lower()
+
+    if ext not in ALLOWED:
+        raise ValueError('File extension should be one of %s, not %s' % (', '.join(ALLOWED), ext))
+
+    try:
+        resp = requests.get(url)
+    except Exception as e:
+        raise Exception('Failed to fetch URL: %s' % e)
+    
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.' + ext) as fp:
+        fp.write(resp.content)
+        return fp.name