units: fix __rfloordiv__ handling of unsupported operands

Author: koreno

easypy/units.py

@@ -72,6 +72,8 @@ def __mul__(self, mul):
     __rmul__ = __mul__
 
     def __floordiv__(self, div):
+        if not isinstance(div, numbers.Number):
+            return NotImplemented
         res = int.__floordiv__(self, div)
         if not isinstance(div, self.__class__):
             # if div is not a datasize, return a datasize
@@ -80,6 +82,8 @@ def __floordiv__(self, div):
         return res
 
     def __rfloordiv__(self, div):
+        if not isinstance(div, numbers.Number):
+            return NotImplemented
         res = int.__rfloordiv__(self, div)
         if not isinstance(div, self.__class__):
             # if div is not a datasize, return a datasize

tests/test_units.py

@@ -1,4 +1,4 @@
-from easypy.units import byte
+from easypy.units import byte, KiB
 
 
 def test_data_sizes():
@@ -10,3 +10,22 @@ def test_data_sizes():
     assert "{0!s}, {0!r}, {0:10text}".format(2**21 * byte) == "2MiB, 2*MiB,       2MiB"
     assert "{0!s}, {0!r}, {0:10text}".format(2**21 * byte + 100) == "~2.0MiB, 2097252*bytes,    ~2.0MiB"
     assert "{0!s}, {0!r}, {0:10text}".format(2**41 * byte + 100) == "~2.0TiB, 2199023255652*bytes,    ~2.0TiB"
+
+
+def test_operators():
+
+    assert (byte * 1024) == KiB
+    assert KiB / 1024 == byte
+    assert KiB / KiB == 1
+    assert KiB / 7 == 146.28571428571428
+    assert KiB // 7 == 146
+    assert 2050 // KiB == (2 * byte)
+
+    # check that __r*__ overloads are used when the unit doesn't support the right-hand operand
+    class Foo():
+        def __rfloordiv__(self, div):
+            return self
+
+    foo = Foo()
+
+    assert KiB // foo is foo