Add download for image previews

Author: seignovert

opus/images.py

@@ -1,5 +1,7 @@
 # -*- coding: utf-8 -*-
 
+from .wget import download
+
 class Image(object):
     def __init__(self, ring_obs_id, path, img):
         self.ring_obs_id = ring_obs_id
@@ -16,6 +18,8 @@ def __str__(self):
     def url(self):
         return self.path + self.img
 
+    def download(self, out=None):
+        download(self.url, out)
 
 class Images(object):
     def __init__(self, json, size):

opus/wget.py

@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+import os
+import requests
+import shutil
+
+def download(url, out=None):
+    outdir = None
+    if out and os.path.isdir(out):
+        outdir = out
+        out = None
+
+    if out is None:
+        filename = url.split('/')[-1]
+    else:
+        filename = out
+
+    if outdir:
+        filename = os.path.join(outdir, filename)
+
+    r = requests.get(url, stream=True)
+    with open(filename, 'wb') as f:
+        shutil.copyfileobj(r.raw, f)
+
+    return filename

requirements.txt

@@ -1,4 +1,3 @@
-wget
 requests
 six
 argparse

tests/api/image/med/N1459551972_1_med.jpg

@@ -1,25 +1,25 @@
 # -*- coding: utf-8 -*-
 import pytest
+import os
 import responses
 import json as JSON
 
 from opus.images import *
 
-
 @pytest.fixture
 def size():
     return 'med'
 
-
-@pytest.fixture
-def json(size):
-    return JSON.loads(open('tests/api/images/{}.json'.format(size), 'r').read())
-
-
 @pytest.fixture
-def images(json, size):
+def images(size):
+    json = JSON.loads(open('tests/api/images/{}.json'.format(size), 'r').read())
     return Images(json, size)
 
+@pytest.fixture
+def image(size):
+    ring_obs_id = 'S_IMG_CO_ISS_1459551972_N'
+    json = JSON.loads(open('tests/api/image/'+size+'/'+ring_obs_id+'.json', 'r').read())
+    return Image(ring_obs_id, json['path'], json['data'][0]['img'])
 
 def test_images_meta(images, size):
     assert repr(images) == 'OPUS API Images object (with 10 images)'
@@ -42,3 +42,43 @@ def test_images_iter(images):
     images.next() == images[0]
     for ii, img in enumerate(images):
         img = images[ii]
+
+
+@responses.activate
+def test_image_download(image):
+    fname = os.path.basename(image.url)
+
+    with open('tests/api/image/med/'+fname, 'rb') as img:
+        responses.add(responses.GET, image.url,
+                    body=img.read(), status=200,
+                    content_type='image/jpeg',
+                    stream=True
+        )
+
+    image.download(out='tests/')
+
+    assert len(responses.calls) == 1
+    assert responses.calls[0].request.url == 'https://pds-rings.seti.org/holdings/previews/COISS_2xxx/COISS_2001/data/1459551663_1459568594/N1459551972_1_med.jpg'
+
+    assert os.path.isfile('tests/'+fname)
+    os.remove('tests/'+fname)
+
+
+@responses.activate
+def test_image_download_output(image):
+    fname = os.path.basename(image.url)
+
+    with open('tests/api/image/med/'+fname, 'rb') as img:
+        responses.add(responses.GET, image.url,
+                      body=img.read(), status=200,
+                      content_type='image/jpeg',
+                      stream=True
+                      )
+
+    image.download(out='tests/test.jpg')
+
+    assert len(responses.calls) == 1
+    assert responses.calls[0].request.url == 'https://pds-rings.seti.org/holdings/previews/COISS_2xxx/COISS_2001/data/1459551663_1459568594/N1459551972_1_med.jpg'
+
+    assert os.path.isfile('tests/test.jpg')
+    os.remove('tests/test.jpg')