From 18107bfc22cb1ee35504750826e7d493b4328110 Mon Sep 17 00:00:00 2001
From: Adam Williamson <awilliam@redhat.com>
Date: Aug 28 2015 07:21:44 +0000
Subject: add size calculation for images


images now have a 'size' property. Yay sizes.

---

diff --git a/fedfind/helpers.py b/fedfind/helpers.py
index 3920b68..4f41f9f 100644
--- a/fedfind/helpers.py
+++ b/fedfind/helpers.py
@@ -124,6 +124,28 @@ def url_exists(url):
     else:
         raise ValueError("Invalid or unhandled URL!")
 
+def get_size(url):
+    """Get the size of URL (Content-Length header for HTTP). Currently
+    only supports HTTP.
+    """
+    logger.debug("Checking size of %s", url)
+    if not url.startswith('http'):
+        raise ValueError("Can only check HTTP URLs for now.")
+    tries = 3
+    headers = None
+    while not headers and tries:
+        try:
+            headers = urlopen(url).info()
+        except (ValueError, URLError, HTTPError):
+            logger.debug("HTTP error! Retrying...")
+            tries -= 1
+    if not headers:
+        raise ValueError("URL {0} not valid or not reachable!".format(url))
+    size = headers.getheader('Content-Length')
+    if not size:
+        raise ValueError("Size could not be found!")
+    return int(size)
+
 def comma_list(string):
     """Split a comma-separated list of values into a list (used by the
     CLI for passing query parameters etc).
diff --git a/fedfind/image.py b/fedfind/image.py
index ecc234a..c0e0eec 100644
--- a/fedfind/image.py
+++ b/fedfind/image.py
@@ -256,3 +256,15 @@ class Image(object):
             rsync = '{0}/{1}'.format(fedfind.const.RSYNC, self.path)
             return fedfind.helpers.url_exists(rsync)
         return fedfind.helpers.url_exists(self.url)
+
+    @property
+    def size(self):
+        """Size of the image (in bytes)."""
+        # This is actually not safe because self.url could be rsync,
+        # but in practice at present it never is. I should fix that
+        # though.
+        try:
+            # Using the mirror system for this seems error-prone.
+            return fedfind.helpers.get_size(self.direct_url)
+        except AttributeError:
+            return fedfind.helpers.get_size(self.url)