util/disk: extend Size class

The Size class can now parse strings and has support for arithmetic
operations and comparisons with numbers.

Signed-off-by: Jan Fajerski <jfajerski@suse.com>
(cherry picked from commit dd89f467a0b89a436371c314bc69cf71a6e9ce72)

diff --git a/src/ceph-volume/ceph_volume/tests/util/test_disk.py b/src/ceph-volume/ceph_volume/tests/util/test_disk.py
index 4f776567d336c3770cd602807ae8beb709dd7b23..5f4d57343e4c2516e2ff9d540bfa59285fa441d0 100644 (file)
--- a/src/ceph-volume/ceph_volume/tests/util/test_disk.py
+++ b/src/ceph-volume/ceph_volume/tests/util/test_disk.py
@@ -164,6 +164,49 @@ class TestSizeFromHumanReadable(object):
         assert result == disk.Size(tb=1.8)
 
 
+class TestSizeParse(object):
+
+    def test_bytes(self):
+        result = disk.Size.parse('2')
+        assert result == disk.Size(b=2)
+
+    def test_kilobytes(self):
+        result = disk.Size.parse('2K')
+        assert result == disk.Size(kb=2)
+
+    def test_megabytes(self):
+        result = disk.Size.parse('2M')
+        assert result == disk.Size(mb=2)
+
+    def test_gigabytes(self):
+        result = disk.Size.parse('2G')
+        assert result == disk.Size(gb=2)
+
+    def test_terrabytes(self):
+        result = disk.Size.parse('2T')
+        assert result == disk.Size(tb=2)
+
+    def test_tb(self):
+        result = disk.Size.parse('2Tb')
+        assert result == disk.Size(tb=2)
+
+    def test_case(self):
+        result = disk.Size.parse('2t')
+        assert result == disk.Size(tb=2)
+
+    def test_space(self):
+        result = disk.Size.parse('2T')
+        assert result == disk.Size(tb=2)
+
+    def test_float(self):
+        result = disk.Size.parse('2.0')
+        assert result == disk.Size(b=2)
+        result = disk.Size.parse('2.0T')
+        assert result == disk.Size(tb=2)
+        result = disk.Size.parse('1.8T')
+        assert result == disk.Size(tb=1.8)
+
+
 class TestGetBlockDevsLsblk(object):
 
     @patch('ceph_volume.process.call')

diff --git a/src/ceph-volume/ceph_volume/util/disk.py b/src/ceph-volume/ceph_volume/util/disk.py
index 2b8a1a9312f354817b18bb5f3552de6bd3eacd20..9874a2dc35850d731a91e7eafbb82904b28af63f 100644 (file)
--- a/src/ceph-volume/ceph_volume/util/disk.py
+++ b/src/ceph-volume/ceph_volume/util/disk.py
@@ -454,15 +454,27 @@ class Size(object):
         Total size: 2.16 GB
     """
 
+    @classmethod
+    def parse(cls, size):
+        if (len(size) > 2 and
+            size[-2].lower() in ['k', 'm', 'g', 't'] and
+            size[-1].lower() == 'b'):
+            return cls(**{size[-2:].lower(): float(size[0:-2])})
+        elif size[-1].lower() in ['b', 'k', 'm', 'g', 't']:
+            return cls(**{size[-1].lower(): float(size[0:-1])})
+        else:
+            return cls(b=float(size))
+
+
     def __init__(self, multiplier=1024, **kw):
         self._multiplier = multiplier
         # create a mapping of units-to-multiplier, skip bytes as that is
         # calculated initially always and does not need to convert
         aliases = [
-            [('kb', 'kilobytes'), self._multiplier],
-            [('mb', 'megabytes'), self._multiplier ** 2],
-            [('gb', 'gigabytes'), self._multiplier ** 3],
-            [('tb', 'terabytes'), self._multiplier ** 4],
+            [('k', 'kb', 'kilobytes'), self._multiplier],
+            [('m', 'mb', 'megabytes'), self._multiplier ** 2],
+            [('g', 'gb', 'gigabytes'), self._multiplier ** 3],
+            [('t', 'tb', 'terabytes'), self._multiplier ** 4],
         ]
         # and mappings for units-to-formatters, including bytes and aliases for
         # each
@@ -519,23 +531,47 @@ class Size(object):
     def __format__(self, spec):
         return str(self._get_best_format()).__format__(spec)
 
+    def __int__(self):
+        return int(self._b)
+
+    def __float__(self):
+        return self._b
+
     def __lt__(self, other):
-        return self._b < other._b
+        if isinstance(other, Size):
+            return self._b < other._b
+        else:
+            return self.b < other
 
     def __le__(self, other):
-        return self._b <= other._b
+        if isinstance(other, Size):
+            return self._b <= other._b
+        else:
+            return self.b <= other
 
     def __eq__(self, other):
-        return self._b == other._b
+        if isinstance(other, Size):
+            return self._b == other._b
+        else:
+            return self.b == other
 
     def __ne__(self, other):
-        return self._b != other._b
+        if isinstance(other, Size):
+            return self._b != other._b
+        else:
+            return self.b != other
 
     def __ge__(self, other):
-        return self._b >= other._b
+        if isinstance(other, Size):
+            return self._b >= other._b
+        else:
+            return self.b >= other
 
     def __gt__(self, other):
-        return self._b > other._b
+        if isinstance(other, Size):
+            return self._b > other._b
+        else:
+            return self.b > other
 
     def __add__(self, other):
         if isinstance(other, Size):