diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index 95c8c448d642..f47dcb52ceb7 100644 --- a/mypyc/irbuild/prepare.py +++ b/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 diff --git a/mypyc/irbuild/util.py b/mypyc/irbuild/util.py index 757b49c68c83..eca2cac7e9db 100644 --- a/mypyc/irbuild/util.py +++ b/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,9 +112,10 @@ 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: @@ -120,10 +123,12 @@ def get_mypyc_attrs(stmt: ClassDef | Decorator) -> dict[str, Any]: 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: diff --git a/mypyc/test-data/irbuild-classes.test b/mypyc/test-data/irbuild-classes.test index bb55958dc6dc..2f59bb000220 100644 --- a/mypyc/test-data/irbuild-classes.test +++ b/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 diff --git a/mypyc/test-data/run-classes.test b/mypyc/test-data/run-classes.test index b25dc9458fd1..79ad2fa0a03b 100644 --- a/mypyc/test-data/run-classes.test +++ b/mypyc/test-data/run-classes.test @@ -3794,3 +3794,39 @@ assert t.native == 43 assert t.generic == "{}" assert t.bitfield == 0x0C assert t.default == 10 + +[case testPerTypeFreeList] +from __future__ import annotations + +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