aboutsummaryrefslogtreecommitdiffstats
path: root/gsxws
diff options
context:
space:
mode:
Diffstat (limited to 'gsxws')
-rw-r--r--gsxws/parts.py14
-rw-r--r--gsxws/products.py15
-rw-r--r--gsxws/utils.py31
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