[mypyc] Allow defining a single-item free "list" for a native class

mypyc/irbuild/prepare.py

@@ -351,13 +351,23 @@ def prepare_class_def(
     ir = mapper.type_to_ir[cdef.info]
     info = cdef.info
 
-    attrs = get_mypyc_attrs(cdef)
+    attrs, attrs_lines = get_mypyc_attrs(cdef)
     if attrs.get("allow_interpreted_subclasses") is True:
         ir.allow_interpreted_subclasses = True
     if attrs.get("serializable") is True:
         # Supports copy.copy and pickle (including subclasses)
         ir._serializable = True
 
+    free_list_len = attrs.get("free_list_len")
+    if free_list_len is not None:
+        line = attrs_lines["free_list_len"]
+        if ir.is_trait:
+            errors.error('"free_list_len" can\'t be used with traits', path, line)
+        if free_list_len == 1:
+            ir.reuse_freed_instance = True
+        else:
+            errors.error(f'Unsupported value for "free_list_len": {free_list_len}', path, line)
+
     # Check for subclassing from builtin types
     for cls in info.mro:
         # Special case exceptions and dicts

mypyc/irbuild/util.py

@@ -96,6 +96,8 @@ def get_mypyc_attr_literal(e: Expression) -> Any:
         return False
     elif isinstance(e, RefExpr) and e.fullname == "builtins.None":
         return None
+    elif isinstance(e, IntExpr):
+        return e.value
     return NotImplemented
 
 
@@ -110,20 +112,23 @@ def get_mypyc_attr_call(d: Expression) -> CallExpr | None:
     return None
 
 
-def get_mypyc_attrs(stmt: ClassDef | Decorator) -> dict[str, Any]:
+def get_mypyc_attrs(stmt: ClassDef | Decorator) -> tuple[dict[str, Any], dict[str, int]]:
     """Collect all the mypyc_attr attributes on a class definition or a function."""
     attrs: dict[str, Any] = {}
+    lines: dict[str, int] = {}
     for dec in stmt.decorators:
         d = get_mypyc_attr_call(dec)
         if d:
             for name, arg in zip(d.arg_names, d.args):
                 if name is None:
                     if isinstance(arg, StrExpr):
                         attrs[arg.value] = True
+                        lines[arg.value] = d.line
                 else:
                     attrs[name] = get_mypyc_attr_literal(arg)
+                    lines[name] = d.line
 
-    return attrs
+    return attrs, lines
 
 
 def is_extension_class(path: str, cdef: ClassDef, errors: Errors) -> bool:

mypyc/test-data/irbuild-classes.test

@@ -1734,3 +1734,23 @@ class NonNative:
 class InheritsPython(dict):
     def __new__(cls) -> InheritsPython:
         return super().__new__(cls)  # E: super().__new__() not supported for classes inheriting from non-native classes
+
+[case testClassWithFreeList]
+from mypy_extensions import mypyc_attr, trait
+
+@mypyc_attr(free_list_len=1)
+class UsesFreeList:
+    pass
+
+@mypyc_attr(free_list_len=None)
+class NoFreeList:
+    pass
+
+@mypyc_attr(free_list_len=2)  # E: Unsupported value for "free_list_len": 2
+class FreeListError:
+    pass
+
+@trait
+@mypyc_attr(free_list_len=1)  # E: "free_list_len" can't be used with traits
+class NonNative:
+    pass

mypyc/test-data/run-classes.test

@@ -3794,3 +3794,37 @@ assert t.native == 43
 assert t.generic == "{}"
 assert t.bitfield == 0x0C
 assert t.default == 10
+
+[case testPerTypeFreeList]
+from mypy_extensions import mypyc_attr
+
+a = []
+
+@mypyc_attr(free_list=1)
+class Foo:
+    def __init__(self, x: int) -> None:
+        self.x = x
+        a.append(x)
+
+def test_alloc() -> None:
+    x: Foo | None
+    y: Foo | None
+
+    x = Foo(1)
+    assert x.x == 1
+    x = None
+
+    x = Foo(2)
+    assert x.x == 2
+    y = Foo(3)
+    assert x.x == 2
+    assert y.x == 3
+    x = None
+    y = None
+    assert a == [1, 2, 3]
+
+    x = Foo(4)
+    assert x.x == 4
+    y = Foo(5)
+    assert x.x == 4
+    assert y.x == 5