Merge pull request #671 from labrys/convert-size

Add convert size function and tests

sickrage/helper/common.py

@@ -157,10 +157,15 @@ def is_torrent_or_nzb_file(filename):
     return filename.rpartition('.')[2].lower() in ['nzb', 'torrent']
 
 
-def pretty_file_size(size):
+def pretty_file_size(size, use_decimal=False, **kwargs):
     """
     Return a human readable representation of the provided ``size``.
+
     :param size: The size to convert
+    :param use_decimal: use decimal instead of binary prefixes (e.g. kilo = 1000 instead of 1024)
+
+    :keyword units: A list of unit names in ascending order. Default units: [u'B', u'KB', u'MB', u'GB', u'TB', u'PB']
+
     :return: The converted size
     """
     try:
@@ -169,13 +174,58 @@ def pretty_file_size(size):
         size = 0.
 
     remaining_size = size
-    for unit in ['B', 'KB', 'MB', 'GB', 'TB', 'PB']:
-        if remaining_size < 1024.:
+    units = kwargs.pop('units', ['B', 'KB', 'MB', 'GB', 'TB', 'PB'])
+    block = 1024. if not use_decimal else 1000.
+    for unit in units:
+        if remaining_size < block:
             return '%3.2f %s' % (remaining_size, unit)
-        remaining_size /= 1024.
+        remaining_size /= block
     return size
 
 
+def convert_size(size, default=None, use_decimal=False, **kwargs):
+    """
+    Convert a file size into the number of bytes
+
+    :param size: to be converted
+    :param default: value to return if conversion fails
+    :param use_decimal: use decimal instead of binary prefixes (e.g. kilo = 1000 instead of 1024)
+
+    :keyword units: A list of (uppercase) unit names in ascending order. Default units: [u'B', u'KB', u'MB', u'GB', u'TB', u'PB']
+
+    :returns: the number of bytes, the default value, or 0
+    """
+    result = None
+
+    try:
+        size_tuple = size.strip().split(u' ')
+        scalar, units = size_tuple[0], size_tuple[1:]
+
+        scalar = float(scalar)
+        units = units[0].upper() if units else None
+
+        scale = kwargs.pop(u'units', [u'B', u'KB', u'MB', u'GB', u'TB', u'PB'])
+        scalar *= (1024 if not use_decimal else 1000) ** scale.index(units)
+
+        result = scalar
+
+    except AttributeError:
+        result = size if size is not None else default
+
+    except ValueError:
+        result = default
+
+    finally:
+        try:
+            result = long(result)
+            if result != default:
+                result = max(result, 0)
+        except (TypeError, ValueError):
+            pass
+
+    return result
+
+
 def remove_extension(filename):
     """
     Remove the extension of the provided ``filename``.
@@ -238,5 +288,5 @@ def try_int(candidate, default_value=0):
 
     try:
         return int(candidate)
-    except Exception:
+    except (ValueError, TypeError):
         return default_value

tests/sickrage_tests/helper/common_tests.py

@@ -34,7 +34,7 @@ sys.path.insert(1, os.path.abspath(os.path.join(os.path.dirname(__file__), '../.
 
 import sickbeard
 from sickrage.helper.common import http_code_description, is_sync_file, is_torrent_or_nzb_file, pretty_file_size
-from sickrage.helper.common import remove_extension, replace_extension, sanitize_filename, try_int
+from sickrage.helper.common import remove_extension, replace_extension, sanitize_filename, try_int, convert_size
 
 
 class CommonTests(unittest.TestCase):
@@ -397,6 +397,21 @@ class CommonTests(unittest.TestCase):
             for (candidate, result) in test.iteritems():
                 self.assertEqual(try_int(candidate, default_value), result)
 
+    def test_convert_size(self):
+        self.assertEqual(convert_size('1 B'), 1)
+        self.assertEqual(convert_size('1 KB'), 1024)
+        self.assertEqual(convert_size('elephant', 'frog'), 'frog')
+        self.assertEqual(convert_size(100, -1), 100)
+        self.assertEqual(convert_size(1.312, -1), 1)
+        self.assertEqual(convert_size(None, -1), -1)
+        self.assertEqual(convert_size(0, -1), 0)
+        self.assertEqual(convert_size(None) or -1, -1)
+        self.assertEqual(convert_size('1 kb', use_decimal=True), 1000)  # Wrong units so result is None
+
+        french_units = ['O', 'KO', 'MO', 'GO', 'TO', 'PO']
+        self.assertEqual(convert_size('1 o', units=french_units), 1)
+        self.assertEqual(convert_size('1 o'), None)  # Wrong units so result is None
+        self.assertEqual(convert_size('1 go', use_decimal=True, units=french_units), 1000000000)  # Wrong units so result is None
 
 if __name__ == '__main__':
     print('=====> Testing %s' % __file__)