diff core/download.py @ 972:7138883966b3

Started unit tests for hotlinking images.
author Brian Neal <bgneal@gmail.com>
date Wed, 23 Sep 2015 21:26:09 -0500
parents
children 6f55c086db1e
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/core/download.py	Wed Sep 23 21:26:09 2015 -0500
@@ -0,0 +1,61 @@
+"""This module contains routines for downloading files."""
+
+import logging
+import os
+import shutil
+import tempfile
+
+import requests
+
+
+logger = logging.getLogger(__name__)
+
+
+def download_file(url, path=None):
+    """Downloads the image file from the given source URL and stores it in the
+    filename given by path. If path is None, a temporary file will be created.
+
+    If successful returns the path to the downloaded file. Otherwise None is
+    returned.
+
+    This function may raise various exceptions from the requests library.
+    """
+    logger.info("download_file from %s; path=%s", url, path)
+
+    try:
+        r = requests.get(url, stream=True)
+    except requests.RequestException:
+        logger.exception("download_file requests.get('%s') exception", url)
+        raise
+
+    if r.status_code != 200:
+        logger.error("download_file from %s: error code %d", url, r.status_code)
+        return None
+
+    # Save file data
+
+    if not path:
+        fd, path = tempfile.mkstemp()
+        os.close(fd)
+
+    try:
+        with open(path, 'wb') as fp:
+            r.raw.decode_content = True
+            shutil.copyfileobj(r.raw, fp)
+    except requests.RequestException:
+        logger.exception("download_file download exception")
+        raise
+
+    file_size = os.stat(path).st_size
+    logger.info("download_file retrieved %s bytes from %s; saved to %s", file_size, url, path)
+    return path
+
+
+if __name__ == '__main__':
+    import sys
+    s = "%(asctime)s : %(levelname)s : %(message)s"
+    logging.basicConfig(level=logging.DEBUG, format=s)
+    logging.info("argument is %s", sys.argv[1])
+    result = download_file(sys.argv[1])
+    if result:
+        print result