improve: preserve Decimal type when round-tripping

hukkin · Jan 9, 2025 · 9c77cbb · 9c77cbb
1 parent e11e00e
commit 9c77cbb
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/src/tomli_w/_writer.py b/src/tomli_w/_writer.py
@@ -125,7 +125,8 @@ def format_decimal(obj: Decimal) -> str:
         return "inf"
     if obj == Decimal("-inf"):
         return "-inf"
-    return str(obj)
+    dec_str = str(obj)
+    return dec_str if "." in dec_str or "E" in dec_str else dec_str + ".0"
 
 
 def format_inline_table(obj: Mapping, ctx: Context) -> str:

diff --git a/tests/test_types.py b/tests/test_types.py
@@ -5,20 +5,26 @@
 
 def test_decimal():
     obj = {
-        "decimal-0": Decimal(0),
+        "decimal-0": Decimal("0"),
+        "decimal-4": Decimal("4"),
         "decimal-pi": Decimal("3.14159"),
         "decimal-inf": Decimal("inf"),
         "decimal-minus-inf": Decimal("-inf"),
         "decimal-nan": Decimal("nan"),
+        "decimal-2e3": Decimal("2e3"),
+        "decimal-2E3": Decimal("2E3"),
     }
     assert (
         tomli_w.dumps(obj)
         == """\
-decimal-0 = 0
+decimal-0 = 0.0
+decimal-4 = 4.0
 decimal-pi = 3.14159
 decimal-inf = inf
 decimal-minus-inf = -inf
 decimal-nan = nan
+decimal-2e3 = 2E+3
+decimal-2E3 = 2E+3
 """
     )