[mypyc] Initial support for annotating classes to be non extension

Author: svalentin

mypyc/irbuild/util.py

@@ -126,14 +126,19 @@ def get_mypyc_attrs(stmt: ClassDef | Decorator) -> dict[str, Any]:
 
 
 def is_extension_class(cdef: ClassDef) -> bool:
-    if any(
-        not is_trait_decorator(d)
-        and not is_dataclass_decorator(d)
-        and not get_mypyc_attr_call(d)
-        and not is_final_decorator(d)
-        for d in cdef.decorators
-    ):
-        return False
+    for d in cdef.decorators:
+        mypyc_attr_call = get_mypyc_attr_call(d)
+        # Classes decorated with "@mypyc_attr(non_extension_class=True)" are not extension classes
+        if mypyc_attr_call and "non_extension_class" in mypyc_attr_call.arg_names:
+            return False
+
+        # Classes that have any decorator other than supported decorators, are not extension classes
+        if (not is_trait_decorator(d)
+            and not is_dataclass_decorator(d)
+            and not mypyc_attr_call
+            and not is_final_decorator(d)):
+            return False
+
     if cdef.info.typeddict_type:
         return False
     if cdef.info.is_named_tuple:

mypyc/test-data/fixtures/ir.py

@@ -350,8 +350,10 @@ def next(i: Iterator[_T]) -> _T: pass
 def next(i: Iterator[_T], default: _T) -> _T: pass
 def hash(o: object) -> int: ...
 def globals() -> Dict[str, Any]: ...
+def hasattr(obj: object, name: str) -> bool: ...
 def getattr(obj: object, name: str, default: Any = None) -> Any: ...
 def setattr(obj: object, name: str, value: Any) -> None: ...
+def delattr(obj: object, name: str) -> None: ...
 def enumerate(x: Iterable[_T]) -> Iterator[Tuple[int, _T]]: ...
 @overload
 def zip(x: Iterable[_T], y: Iterable[_S]) -> Iterator[Tuple[_T, _S]]: ...

mypyc/test-data/run-classes.test

@@ -2829,3 +2829,40 @@ Traceback (most recent call last):
   File "native.py", line 5, in __del__
     raise Exception("e2")
 Exception: e2
+
+[case testNonExtensionClassAttr]
+from mypy_extensions import mypyc_attr
+
+@mypyc_attr(non_extension_class=True)
+class AnnontatedNonExtensionClass:
+    pass
+
+class DerivedClass(AnnontatedNonExtensionClass):
+    pass
+
+def test_function():
+    setattr(AnnontatedNonExtensionClass, 'attr_class', 5)
+    assert(hasattr(AnnontatedNonExtensionClass, 'attr_class') == True)
+    assert(getattr(AnnontatedNonExtensionClass, 'attr_class') == 5)
+    delattr(AnnontatedNonExtensionClass, 'attr_class')
+    assert(hasattr(AnnontatedNonExtensionClass, 'attr_class') == False)
+
+    inst = AnnontatedNonExtensionClass()
+    setattr(inst, 'attr_instance', 6)
+    assert(hasattr(inst, 'attr_instance') == True)
+    assert(getattr(inst, 'attr_instance') == 6)
+    delattr(inst, 'attr_instance')
+    assert(hasattr(inst, 'attr_instance') == False)
+
+    setattr(DerivedClass, 'attr_class', 5)
+    assert(hasattr(DerivedClass, 'attr_class') == True)
+    assert(getattr(DerivedClass, 'attr_class') == 5)
+    delattr(DerivedClass, 'attr_class')
+    assert(hasattr(DerivedClass, 'attr_class') == False)
+
+    derived_inst = DerivedClass()
+    setattr(derived_inst, 'attr_instance', 6)
+    assert(hasattr(derived_inst, 'attr_instance') == True)
+    assert(getattr(derived_inst, 'attr_instance') == 6)
+    delattr(derived_inst, 'attr_instance')
+    assert(hasattr(derived_inst, 'attr_instance') == False)