fix: Fix division for integers (#312)

python-poetry · Sep 20, 2023 · 1b2a53f · 1b2a53f
1 parent 7ed7d3a
commit 1b2a53f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/tests/test_items.py b/tests/test_items.py
@@ -624,6 +624,10 @@ def test_integers_behave_like_ints():
     assert i == 33
     assert i.as_string() == "33"
 
+    i /= 2
+    assert i == 16.5
+    assert i.as_string() == "16.5"
+
     doc = parse("int = +34")
     doc["int"] += 1
 

diff --git a/tomlkit/items.py b/tomlkit/items.py
@@ -667,12 +667,22 @@ def _getstate(self, protocol=3):
     __rpow__ = wrap_method(int.__rpow__)
     __rrshift__ = wrap_method(int.__rrshift__)
     __rshift__ = wrap_method(int.__rshift__)
-    __rtruediv__ = wrap_method(int.__rtruediv__)
     __rxor__ = wrap_method(int.__rxor__)
-    __truediv__ = wrap_method(int.__truediv__)
     __trunc__ = wrap_method(int.__trunc__)
     __xor__ = wrap_method(int.__xor__)
 
+    def __rtruediv__(self, other):
+        result = int.__rtruediv__(self, other)
+        if result is NotImplemented:
+            return result
+        return Float._new(self, result)
+
+    def __truediv__(self, other):
+        result = int.__truediv__(self, other)
+        if result is NotImplemented:
+            return result
+        return Float._new(self, result)
+
 
 class Float(Item, _CustomFloat):
     """