[mypyc] Keep Generic bases of non-native generic classes

mypyc/ir/class_ir.py

@@ -622,11 +622,20 @@ class NonExtClassInfo:
     the class annotations dictionary, and the metaclass.
     """
 
-    def __init__(self, dict: Value, bases: Value, anns: Value, metaclass: Value) -> None:
+    def __init__(
+        self,
+        dict: Value,
+        bases: Value,
+        anns: Value,
+        metaclass: Value,
+        orig_bases: Value | None = None,
+    ) -> None:
         self.dict = dict
         self.bases = bases
         self.anns = anns
         self.metaclass = metaclass
+        # Original expressions, before resolving __mro_entries__, if resolution was needed.
+        self.orig_bases = orig_bases
 
 
 def serialize_vtable_entry(entry: VTableMethod) -> JsonDict:

mypyc/irbuild/classdef.py

@@ -17,6 +17,7 @@
     EllipsisExpr,
     ExpressionStmt,
     FuncDef,
+    IndexExpr,
     Lvalue,
     MemberExpr,
     NameExpr,
@@ -80,6 +81,7 @@
     not_implemented_op,
     py_calc_meta_op,
     py_init_subclass_op,
+    py_resolve_bases_op,
     pytype_from_template_op,
     type_object_op,
 )
@@ -265,16 +267,35 @@ def class_body_obj(self) -> Value | None:
         return self.non_ext.dict
 
     def create_non_ext_info(self) -> NonExtClassInfo:
-        non_ext_bases = populate_non_ext_bases(self.builder, self.cdef)
-        non_ext_metaclass = find_non_ext_metaclass(self.builder, self.cdef, non_ext_bases)
-        non_ext_dict = setup_non_ext_dict(
-            self.builder, self.cdef, non_ext_metaclass, non_ext_bases
-        )
+        builder, cdef = self.builder, self.cdef
+        generic_base = None
+        type_params = None
+        if cdef.type_args:
+            # Create type parameters first, since base class expressions can refer to them.
+            generic_base, tvs = make_generic_base_class(
+                builder, cdef.fullname, cdef.type_args, cdef.line
+            )
+            type_params = builder.new_tuple(tvs, cdef.line)
+        non_ext_bases = populate_non_ext_bases(builder, cdef, generic_base)
+        orig_bases = None
+        if uses_non_ext_base_exprs(cdef):
+            orig_bases = non_ext_bases
+            non_ext_bases = builder.call_c(py_resolve_bases_op, [non_ext_bases], cdef.line)
+        non_ext_metaclass = find_non_ext_metaclass(builder, cdef, non_ext_bases)
+        non_ext_dict = setup_non_ext_dict(builder, cdef, non_ext_metaclass, non_ext_bases)
+        if type_params is not None:
+            builder.call_c(
+                exact_dict_set_item_op,
+                [non_ext_dict, builder.load_str("__type_params__"), type_params],
+                cdef.line,
+            )
         # We populate __annotations__ for non-extension classes
         # because dataclasses uses it to determine which attributes to compute on.
         # TODO: Maybe generate more precise types for annotations
-        non_ext_anns = self.builder.call_c(dict_new_op, [], self.cdef.line)
-        return NonExtClassInfo(non_ext_dict, non_ext_bases, non_ext_anns, non_ext_metaclass)
+        non_ext_anns = builder.call_c(dict_new_op, [], cdef.line)
+        return NonExtClassInfo(
+            non_ext_dict, non_ext_bases, non_ext_anns, non_ext_metaclass, orig_bases
+        )
 
     def add_method(self, fdef: FuncDef) -> None:
         handle_non_ext_method(self.builder, self.non_ext, self.cdef, fdef)
@@ -475,7 +496,10 @@ def allocate_class(builder: IRBuilder, cdef: ClassDef) -> Value:
     base_exprs = cdef.base_type_exprs + cdef.removed_base_type_exprs
     new_style_type_args = cdef.type_args
     if new_style_type_args:
-        bases = [make_generic_base_class(builder, cdef.fullname, new_style_type_args, cdef.line)]
+        generic_base, _ = make_generic_base_class(
+            builder, cdef.fullname, new_style_type_args, cdef.line
+        )
+        bases = [generic_base]
     else:
         bases = []
 
@@ -538,8 +562,11 @@ def allocate_class(builder: IRBuilder, cdef: ClassDef) -> Value:
 
 def make_generic_base_class(
     builder: IRBuilder, fullname: str, type_args: list[TypeParam], line: int
-) -> Value:
-    """Construct Generic[...] base class object for a new-style generic class (Python 3.12)."""
+) -> tuple[Value, list[Value]]:
+    """Construct Generic[...] base class object for a new-style generic class (Python 3.12).
+
+    Also return the created type parameter objects.
+    """
     mod = builder.call_c(import_op, [builder.load_str("_typing")], line)
     tvs = create_type_params(builder, mod, type_args, line)
     args = []
@@ -557,7 +584,7 @@ def make_generic_base_class(
         arg = builder.new_tuple(args, line)
 
     base = builder.primitive_op(py_get_item_op, [gent, arg], line)
-    return base
+    return base, tvs
 
 
 # Mypy uses these internally as base classes of TypedDict classes. These are
@@ -568,13 +595,51 @@ def make_generic_base_class(
 )
 
 
-def populate_non_ext_bases(builder: IRBuilder, cdef: ClassDef) -> Value:
+def uses_non_ext_base_exprs(cdef: ClassDef) -> bool:
+    """Should the base classes of a non-extension class be evaluated as written?
+
+    Resolved base types erase Generic[...] bases and type arguments, but they are
+    needed at runtime, e.g. to make the class generic.
+    """
+    info = cdef.info
+    if info.is_named_tuple or info.typeddict_type is not None:
+        # These have synthesized base classes that need special casing.
+        return False
+    return bool(
+        cdef.type_args
+        or cdef.removed_base_type_exprs
+        or any(isinstance(expr, IndexExpr) for expr in cdef.base_type_exprs)
+    )
+
+
+def populate_non_ext_bases(
+    builder: IRBuilder, cdef: ClassDef, generic_base: Value | None
+) -> Value:
     """Create base class tuple of a non-extension class.
 
-    The tuple is passed to the metaclass constructor.
+    The tuple is passed to the metaclass constructor, after resolving __mro_entries__
+    if uses_non_ext_base_exprs(cdef) is true. If given, generic_base is the implicit
+    Generic[...] base of a new-style generic class.
     """
     is_named_tuple = cdef.info.is_named_tuple
     ir = builder.mapper.type_to_ir[cdef.info]
+    for base_type in cdef.info.bases:
+        # Add the current class to the base classes list of concrete subclasses.
+        base_ir = builder.mapper.type_to_ir.get(base_type.type)
+        if base_ir is not None and base_ir.children is not None:
+            base_ir.children.append(ir)
+
+    if uses_non_ext_base_exprs(cdef):
+        # Mypy removes Generic[...] bases, so restore source order before evaluation.
+        base_exprs = sorted(
+            cdef.base_type_exprs + cdef.removed_base_type_exprs,
+            key=lambda expr: (expr.line, expr.column),
+        )
+        items = [builder.accept(expr) for expr in base_exprs]
+        if generic_base is not None:
+            items.append(generic_base)
+        return builder.new_tuple(items, cdef.line)
+
     bases = []
     for cls in (b.type for b in cdef.info.bases):
         if cls.fullname == "builtins.object":
@@ -590,12 +655,6 @@ def populate_non_ext_bases(builder: IRBuilder, cdef: ClassDef) -> Value:
             # HAX: Synthesized base classes added by mypy don't exist at runtime, so skip them.
             #      This could break if they were added explicitly, though...
             continue
-        # Add the current class to the base classes list of concrete subclasses
-        if cls in builder.mapper.type_to_ir:
-            base_ir = builder.mapper.type_to_ir[cls]
-            if base_ir.children is not None:
-                base_ir.children.append(ir)
-
         if cls.fullname in MAGIC_TYPED_DICT_CLASSES:
             # HAX: Mypy internally represents TypedDict classes differently from what
             #      should happen at runtime. Replace with something that works.
@@ -944,6 +1003,20 @@ def load_non_ext_class(
 
     add_dunders_to_non_ext_dict(builder, non_ext, line)
 
+    if non_ext.orig_bases is not None:
+        changed = builder.builder.translate_is_op(
+            non_ext.bases, non_ext.orig_bases, "is not", line
+        )
+        update, done = BasicBlock(), BasicBlock()
+        builder.add_bool_branch(changed, update, done)
+        builder.activate_block(update)
+        builder.call_c(
+            exact_dict_set_item_op,
+            [non_ext.dict, builder.load_str("__orig_bases__"), non_ext.orig_bases],
+            line,
+        )
+        builder.goto_and_activate(done)
+
     class_type_obj = builder.py_call(
         non_ext.metaclass, [cls_name, non_ext.bases, non_ext.dict], line
     )

mypyc/lib-rt/CPy.h

@@ -1037,6 +1037,7 @@ static inline PyObject *CPy_TYPE(PyObject *obj) {
 }
 
 PyObject *CPy_CalculateMetaclass(PyObject *type, PyObject *o);
+PyObject *CPyType_ResolveBases(PyObject *bases);
 PyObject *CPy_GetCoro(PyObject *obj);
 PyObject *CPyGen_AlreadyExecutingError(int is_coroutine);
 PyObject *CPyIter_Send(PyObject *iter, PyObject *val, PyObject **stop_iter_value);

mypyc/lib-rt/misc_ops.c

@@ -151,6 +151,15 @@ static bool _CPy_IsSafeMetaClass(PyTypeObject *metaclass) {
     return matches;
 }
 
+// Resolve __mro_entries__ in base class expressions, returning an owned reference.
+PyObject *CPyType_ResolveBases(PyObject *bases) {
+    PyObject *resolved = update_bases(bases);
+    if (resolved == bases) {
+        Py_INCREF(resolved);
+    }
+    return resolved;
+}
+
 #if CPY_3_13_FEATURES
 
 // Adapted from CPython 3.13.0b3

mypyc/primitives/misc_ops.py

@@ -131,6 +131,14 @@
     error_kind=ERR_MAGIC,
 )
 
+# Resolve __mro_entries__ in base class expressions before selecting the metaclass.
+py_resolve_bases_op = custom_op(
+    arg_types=[object_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name="CPyType_ResolveBases",
+    error_kind=ERR_MAGIC,
+)
+
 # Determine the most derived metaclass and check for metaclass conflicts.
 # Arguments are (metaclass, bases).
 py_calc_meta_op = custom_op(

mypyc/test-data/run-classes.test

@@ -686,6 +686,79 @@ c.set([1])
 assert c.x == [1]
 assert c.get() == [1]
 
+[case testNonNativeGenericClass]
+from typing import Dict, Generic, TypeVar
+from mypy_extensions import mypyc_attr
+
+T = TypeVar("T")
+S = TypeVar("S")
+
+@mypyc_attr(native_class=False)
+class G(Generic[T]):
+    def __init__(self, x: T) -> None:
+        self.x = x
+
+@mypyc_attr(native_class=False)
+class Known:
+    pass
+
+@mypyc_attr(native_class=False)
+class GenericFirst(Generic[T], Known):
+    pass
+
+@mypyc_attr(native_class=False)
+class GenericLast(Known, Generic[T]):
+    pass
+
+@mypyc_attr(native_class=False)
+class Sub(G[int]):
+    pass
+
+@mypyc_attr(native_class=False)
+class ImplicitGeneric(G[T]):
+    pass
+
+@mypyc_attr(native_class=False)
+class TwoParams(G[T], Generic[S, T]):
+    pass
+
+@mypyc_attr(native_class=False)
+class DictSub(Dict[str, int]):
+    pass
+
+def make() -> G[int]:
+    return G[int](5)
+
+[file driver.py]
+from typing import Generic
+from native import (
+    G, T, S, Known, GenericFirst, GenericLast, Sub, ImplicitGeneric, TwoParams, DictSub, make
+)
+
+assert G.__bases__ == (Generic,)
+assert G.__orig_bases__ == (Generic[T],)
+assert G.__parameters__ == (T,)
+assert G[int](3).x == 3
+assert G[int].__origin__ is G
+assert make().x == 5
+
+assert GenericFirst.__bases__ == (Generic, Known)
+assert GenericLast.__bases__ == (Known, Generic)
+assert GenericFirst.__parameters__ == GenericLast.__parameters__ == (T,)
+
+assert Sub.__bases__ == (G,)
+assert Sub.__orig_bases__ == (G[int],)
+assert Sub.__parameters__ == ()
+assert Sub(7).x == 7
+
+assert ImplicitGeneric.__parameters__ == (T,)
+assert ImplicitGeneric[str]("a").x == "a"
+assert TwoParams.__bases__ == (G, Generic)
+assert TwoParams.__parameters__ == (S, T)
+
+assert DictSub.__bases__ == (dict, Generic)
+assert DictSub(a=1) == {"a": 1}
+
 [case testSubclass1]
 from typing import Tuple
 

mypyc/test-data/run-python312.test

@@ -243,3 +243,52 @@ def test_param_spec_type_alias() -> None:
         assert str(D[[int, str]]) == "D[[int, str]]"
     assert str(getattr(D, "__value__")) == "typing.Callable[P, int]"
 [typing fixtures/typing-full.pyi]
+
+[case testPEP695NonNativeClass]
+from mypy_extensions import mypyc_attr
+
+@mypyc_attr(native_class=False)
+class Base:
+    pass
+
+@mypyc_attr(native_class=False)
+class G[T]:
+    def __init__(self, x: T) -> None:
+        self.x = x
+
+@mypyc_attr(native_class=False)
+class WithBase[T](Base):
+    pass
+
+@mypyc_attr(native_class=False)
+class Sub[S](G[S]):
+    pass
+
+@mypyc_attr(native_class=False)
+class Variadic[T, *Ts, **P]:
+    pass
+
+[file driver.py]
+from typing import Generic
+from native import Base, G, WithBase, Sub, Variadic
+
+(T,) = G.__type_params__
+assert G.__bases__ == (Generic,)
+assert G.__orig_bases__ == (Generic[T],)
+assert G.__parameters__ == (T,)
+assert G[int](3).x == 3
+
+assert WithBase.__bases__ == (Base, Generic)
+assert WithBase.__parameters__ == WithBase.__type_params__
+assert WithBase[int]() is not None
+
+(S,) = Sub.__type_params__
+assert Sub.__bases__ == (G, Generic)
+assert Sub.__orig_bases__ == (G[S], Generic[S])
+assert Sub.__parameters__ == (S,)
+assert Sub[str]("a").x == "a"
+
+assert len(Variadic.__type_params__) == 3
+assert Variadic.__parameters__ == Variadic.__type_params__
+assert Variadic[int, str, [int]] is not None
+[typing fixtures/typing-full.pyi]