add tests

Author: BobTheBuidler

mypyc/codegen/emit.py

@@ -210,7 +210,7 @@ def object_annotation(self, obj: object, line: str) -> str:
 
         If it contains illegal characters, an empty string is returned."""
         line_width = self._indent + len(line)
-        formatted = pformat_deterministic(obj, line_width)
+        formatted = pformat_deterministic(obj, max(90 - line_width, 20))
 
         if any(x in formatted for x in ("/*", "*/", "\0")):
             return ""
@@ -1274,14 +1274,14 @@ def native_function_doc_initializer(func: FuncIR) -> str:
     return c_string_initializer(docstring.encode("ascii", errors="backslashreplace"))
 
 
-def pformat_deterministic(obj: object, line_width: int) -> str:
+def pformat_deterministic(obj: object, width: int) -> str:
     """Pretty-print `obj` with deterministic sorting for mypyc literal types."""
-    # Temporarily override pprint._safe_key
+    # Temporarily override pprint._safe_key to get deterministic ordering of containers.
     default_safe_key = pprint._safe_key  # type: ignore [attr-defined]
     pprint._safe_key = _mypyc_safe_key  # type: ignore [attr-defined]
-    
+
     try:
-        return pprint.pformat(obj, compact=True, width=max(90 - line_width, 20))
+        return pprint.pformat(obj, compact=True, width=width)
     finally:
         # Always restore the original key to avoid affecting other pprint users.
         pprint._safe_key = default_safe_key  # type: ignore [attr-defined]

mypyc/test/test_emit.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+import pprint
 import unittest
 
-from mypyc.codegen.emit import Emitter, EmitterContext
+from mypyc.codegen.emit import Emitter, EmitterContext, pformat_deterministic
 from mypyc.common import HAVE_IMMORTAL
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.ops import BasicBlock, Register, Value
@@ -21,6 +22,37 @@
 from mypyc.namegen import NameGenerator
 
 
+class TestPformatDeterministic(unittest.TestCase):
+    def test_frozenset_elements_sorted(self) -> None:
+        fs_small = frozenset({("a", 1)})
+        fs_large = frozenset({("a", 1), ("b", 2)})
+        literal_a = frozenset({fs_large, fs_small})
+        literal_b = frozenset({fs_small, fs_large})
+        expected = "frozenset({frozenset({('a', 1)}), frozenset({('a', 1), ('b', 2)})})"
+
+        assert pformat_deterministic(literal_a, 80) == expected
+        assert pformat_deterministic(literal_b, 80) == expected
+
+    def test_nested_supported_literals(self) -> None:
+        nested_frozen = frozenset({("m", 0), ("n", 1)})
+        item_a = ("outer", 1, nested_frozen)
+        item_b = ("outer", 2, frozenset({("x", 3)}))
+        literal_a = frozenset({item_a, item_b})
+        literal_b = frozenset({item_b, item_a})
+        expected = (
+            "frozenset({('outer', 1, frozenset({('m', 0), ('n', 1)})), "
+            "('outer', 2, frozenset({('x', 3)}))})"
+        )
+
+        assert pformat_deterministic(literal_a, 120) == expected
+        assert pformat_deterministic(literal_b, 120) == expected
+
+    def test_restores_default_safe_key(self) -> None:
+        original_safe_key = pprint._safe_key
+        pformat_deterministic({"key": "value"}, 80)
+        assert pprint._safe_key is original_safe_key
+
+
 class TestEmitter(unittest.TestCase):
     def setUp(self) -> None:
         self.n = Register(int_rprimitive, "n")