Implement ByteSizeParser.

3 files changed, 76 insertions, 4 deletions

diff --git a/obnamlib/__init__.py b/obnamlib/__init__.py
index b0870df4..b1952403 100644
--- a/obnamlib/__init__.py
+++ b/obnamlib/__init__.py
@@ -29,7 +29,7 @@ class Error(Exception):
 CHUNK_SIZE = 4096
 CHUNK_GROUP_SIZE = 16
 
-from sizeparse import ByteSizeParser
+from sizeparse import SizeSyntaxError, UnitNameError, ByteSizeParser
 
 from hooks import Hook, HookManager
 from cfg import Configuration
diff --git a/obnamlib/sizeparse.py b/obnamlib/sizeparse.py
index d0649646..78dd63ac 100644
--- a/obnamlib/sizeparse.py
+++ b/obnamlib/sizeparse.py
@@ -14,15 +14,61 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 
+import re
+
+
+class UnitError(Exception):
+
+    def __str__(self):
+        return self.msg
+
+
+class SizeSyntaxError(UnitError):
+
+    def __init__(self, string):
+        self.msg = '"%s" is not a valid size' % string
+
+
+class UnitNameError(UnitError):
+
+    def __init__(self, string):
+        self.msg = '"%s" is not a valid unit' % string
+
+
 class ByteSizeParser(object):
 
     '''Parse sizes of data in bytes, kilobytes, kibibytes, etc.'''
     
+    pat = re.compile(r'^(?P<size>\d+(\.\d+)?)\s*'
+                     r'(?P<unit>[kmg]?i?b?)?$', re.I)
+    
+    units = {
+        'b': ('B', 1),
+        'kb': ('kB', 1000),
+        'kib': ('KiB', 1024),
+        'mb': ('kB', 1000**2),
+        'mib': ('KiB', 1024**2),
+        'gb': ('GB', 1000**3),
+        'gib': ('GiB', 1024**3),
+    }
+    
     def __init__(self):
-        pass
+        self.set_default_unit('B')
         
     def set_default_unit(self, unit):
-        pass
+        if unit.lower() not in self.units:
+            raise UnitNameError(unit)
+        self.default_unit = unit
         
     def parse(self, string):
-        pass
+        m = self.pat.match(string)
+        if not m:
+            raise SizeSyntaxError(string)
+        size = float(m.group('size'))
+        unit = m.group('unit')
+        if not unit:
+            unit = self.default_unit
+        elif unit.lower() not in self.units:
+            raise UnitNameError(unit)
+        unit_name, factor = self.units[unit.lower()]
+        return size * factor
diff --git a/obnamlib/sizeparse_tests.py b/obnamlib/sizeparse_tests.py
index 85ddcf55..fcafb653 100644
--- a/obnamlib/sizeparse_tests.py
+++ b/obnamlib/sizeparse_tests.py
@@ -55,3 +55,29 @@ class ByteSizeParserTests(unittest.TestCase):
     def test_parses_size_with_gibibyte_unit(self):
         self.assertEqual(self.p.parse('123 GiB'), 123 * 1024**3)
 
+    def test_raises_error_for_empty_string(self):
+        self.assertRaises(obnamlib.SizeSyntaxError, self.p.parse, '')
+
+    def test_raises_error_for_missing_size(self):
+        self.assertRaises(obnamlib.SizeSyntaxError, self.p.parse, 'KiB')
+
+    def test_raises_error_for_bad_unit(self):
+        self.assertRaises(obnamlib.SizeSyntaxError, self.p.parse, '1 km')
+
+    def test_raises_error_for_bad_unit_thats_similar_to_real_one(self):
+        self.assertRaises(obnamlib.UnitNameError, self.p.parse, '1 ib')
+
+    def test_raises_error_for_bad_default_unit(self):
+        self.assertRaises(obnamlib.UnitNameError, 
+                          self.p.set_default_unit, 'km')
+
+    def test_size_syntax_error_includes_input_string(self):
+        text = 'asdf asdf'
+        e = obnamlib.SizeSyntaxError(text)
+        self.assert_(text in str(e), str(e))
+
+    def test_unit_name_error_includes_input_string(self):
+        text = 'asdf asdf'
+        e = obnamlib.UnitNameError(text)
+        self.assert_(text in str(e), str(e))
+