Add support for resolving metaclasses (#14120)

## Summary

I mirrored some of the idioms that @AlexWaygood used in the MRO work.

Closes https://github.com/astral-sh/ruff/issues/14096.

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

crates/red_knot_python_semantic/resources/mdtest/metaclass.md

@@ -0,0 +1,191 @@
+## Default
+
+```py
+class M(type): ...
+
+reveal_type(M.__class__)  # revealed: Literal[type]
+```
+
+## `object`
+
+```py
+reveal_type(object.__class__)  # revealed: Literal[type]
+```
+
+## `type`
+
+```py
+reveal_type(type.__class__)  # revealed: Literal[type]
+```
+
+## Basic
+
+```py
+class A(type): ...
+class B(metaclass=A): ...
+
+reveal_type(B.__class__)  # revealed: Literal[A]
+```
+
+## Invalid metaclass
+
+If a class is a subclass of a class with a custom metaclass, then the subclass will also have that
+metaclass.
+
+```py
+class M: ...
+class A(metaclass=M): ...
+
+# TODO: emit a diagnostic for the invalid metaclass
+reveal_type(A.__class__)  # revealed: Literal[M]
+```
+
+## Linear inheritance
+
+If a class is a subclass of a class with a custom metaclass, then the subclass will also have that
+metaclass.
+
+```py
+class M(type): ...
+class A(metaclass=M): ...
+class B(A): ...
+
+reveal_type(B.__class__)  # revealed: Literal[M]
+```
+
+## Conflict (1)
+
+The metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its
+bases. ("Strict subclass" is a synonym for "proper subclass"; a non-strict subclass can be a
+subclass or the class itself.)
+
+```py
+class M1(type): ...
+class M2(type): ...
+class A(metaclass=M1): ...
+class B(metaclass=M2): ...
+
+# error: [conflicting-metaclass] "The metaclass of a derived class (`C`) must be a subclass of the metaclasses of all its bases, but `M1` and `M2` have no subclass relationship"
+class C(A, B): ...
+
+reveal_type(C.__class__)  # revealed: Unknown
+```
+
+## Conflict (2)
+
+The metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all its
+bases. ("Strict subclass" is a synonym for "proper subclass"; a non-strict subclass can be a
+subclass or the class itself.)
+
+```py
+class M1(type): ...
+class M2(type): ...
+class A(metaclass=M1): ...
+
+# error: [conflicting-metaclass] "The metaclass of a derived class (`B`) must be a subclass of the metaclasses of all its bases, but `M2` and `M1` have no subclass relationship"
+class B(A, metaclass=M2): ...
+
+reveal_type(B.__class__)  # revealed: Unknown
+```
+
+## Common metaclass
+
+A class has two explicit bases, both of which have the same metaclass.
+
+```py
+class M(type): ...
+class A(metaclass=M): ...
+class B(metaclass=M): ...
+class C(A, B): ...
+
+reveal_type(C.__class__)  # revealed: Literal[M]
+```
+
+## Metaclass metaclass
+
+A class has an explicit base with a custom metaclass. That metaclass itself has a custom metaclass.
+
+```py
+class M1(type): ...
+class M2(type, metaclass=M1): ...
+class M3(M2): ...
+class A(metaclass=M3): ...
+class B(A): ...
+
+reveal_type(A.__class__)  # revealed: Literal[M3]
+```
+
+## Diamond inheritance
+
+```py
+class M(type): ...
+class M1(M): ...
+class M2(M): ...
+class M12(M1, M2): ...
+class A(metaclass=M1): ...
+class B(metaclass=M2): ...
+class C(metaclass=M12): ...
+
+# error: [conflicting-metaclass] "The metaclass of a derived class (`D`) must be a subclass of the metaclasses of all its bases, but `M1` and `M2` have no subclass relationship"
+class D(A, B, C): ...
+
+reveal_type(D.__class__)  # revealed: Unknown
+```
+
+## Unknown
+
+```py
+from nonexistent_module import UnknownClass  # error: [unresolved-import]
+
+class C(UnknownClass): ...
+
+reveal_type(C.__class__)  # revealed: Literal[type]
+
+class M(type): ...
+class A(metaclass=M): ...
+class B(A, UnknownClass): ...
+
+# TODO: This should resolve to `type[M] | Unknown` instead
+reveal_type(B.__class__)  # revealed: Literal[M]
+```
+
+## Duplicate
+
+```py
+class M(type): ...
+class A(metaclass=M): ...
+class B(A, A): ...  # error: [duplicate-base] "Duplicate base class `A`"
+
+reveal_type(B.__class__)  # revealed: Literal[M]
+```
+
+## Non-class
+
+When a class has an explicit `metaclass` that is not a class, the value should be returned as is.
+
+```py
+class A(metaclass=1): ...
+
+reveal_type(A.__class__)  # revealed: Literal[1]
+```
+
+## Cyclic
+
+Retrieving the metaclass of a cyclically defined class should not cause an infinite loop.
+
+```py path=a.pyi
+class A(B): ...  # error: [cyclic-class-def]
+class B(C): ...  # error: [cyclic-class-def]
+class C(A): ...  # error: [cyclic-class-def]
+
+reveal_type(A.__class__)  # revealed: Unknown
+```
+
+## PEP 695 generic
+
+```py
+class M(type): ...
+class A[T: str](metaclass=M): ...
+
+reveal_type(A.__class__)  # revealed: Literal[M]
+```

crates/red_knot_python_semantic/resources/mdtest/scopes/moduletype_attrs.md

@@ -60,8 +60,7 @@ reveal_type(typing.__init__)  # revealed: Literal[__init__]
 # These come from `builtins.object`, not `types.ModuleType`:
 reveal_type(typing.__eq__)  # revealed: Literal[__eq__]
 
-# TODO: understand properties
+reveal_type(typing.__class__)  # revealed: Literal[type]
-reveal_type(typing.__class__)  # revealed: Literal[__class__]
 
 # TODO: needs support for attribute access on instances, properties and generics;
 # should be `dict[str, Any]`

crates/red_knot_python_semantic/src/types.rs

@@ -1,9 +1,11 @@
-use mro::{ClassBase, Mro, MroError, MroIterator};
+use std::hash::Hash;
+
+use indexmap::IndexSet;
+use itertools::Itertools;
+
 use ruff_db::files::File;
 use ruff_python_ast as ast;
 
-use itertools::Itertools;
-
 use crate::semantic_index::ast_ids::HasScopedAstId;
 use crate::semantic_index::definition::Definition;
 use crate::semantic_index::symbol::{self as symbol, ScopeId, ScopedSymbolId};
@@ -16,6 +18,7 @@ use crate::stdlib::{
 };
 use crate::symbol::{Boundness, Symbol};
 use crate::types::diagnostic::TypeCheckDiagnosticsBuilder;
+use crate::types::mro::{ClassBase, Mro, MroError, MroIterator};
 use crate::types::narrow::narrowing_constraint;
 use crate::{Db, FxOrderSet, HasTy, Module, SemanticModel};
 
@@ -1279,8 +1282,7 @@ impl<'db> Type<'db> {
             Type::FunctionLiteral(_) => KnownClass::FunctionType.to_class(db),
             Type::ModuleLiteral(_) => KnownClass::ModuleType.to_class(db),
             Type::Tuple(_) => KnownClass::Tuple.to_class(db),
-            // TODO not accurate if there's a custom metaclass...
+            Type::ClassLiteral(ClassLiteralType { class }) => class.metaclass(db),
-            Type::ClassLiteral(_) => KnownClass::Type.to_class(db),
             // TODO can we do better here? `type[LiteralString]`?
             Type::StringLiteral(_) | Type::LiteralString => KnownClass::Str.to_class(db),
             // TODO: `type[Any]`?
@@ -2026,6 +2028,113 @@ impl<'db> Class<'db> {
         self.iter_mro(db).contains(&ClassBase::Class(other))
     }
 
+    /// Return the explicit `metaclass` of this class, if one is defined.
+    ///
+    /// ## Note
+    /// Only call this function from queries in the same file or your
+    /// query depends on the AST of another file (bad!).
+    fn explicit_metaclass(self, db: &'db dyn Db) -> Option<Type<'db>> {
+        let class_stmt = self.node(db);
+        let metaclass_node = &class_stmt
+            .arguments
+            .as_ref()?
+            .find_keyword("metaclass")?
+            .value;
+        Some(if class_stmt.type_params.is_some() {
+            // when we have a specialized scope, we'll look up the inference
+            // within that scope
+            let model = SemanticModel::new(db, self.file(db));
+            metaclass_node.ty(&model)
+        } else {
+            // Otherwise, we can do the lookup based on the definition scope
+            let class_definition = semantic_index(db, self.file(db)).definition(class_stmt);
+            definition_expression_ty(db, class_definition, metaclass_node)
+        })
+    }
+
+    /// Return the metaclass of this class, or `Unknown` if the metaclass cannot be inferred.
+    pub(crate) fn metaclass(self, db: &'db dyn Db) -> Type<'db> {
+        // TODO: `type[Unknown]` would be a more precise fallback
+        // (needs support for <https://docs.python.org/3/library/typing.html#the-type-of-class-objects>)
+        self.try_metaclass(db).unwrap_or(Type::Unknown)
+    }
+
+    /// Return the metaclass of this class, or an error if the metaclass cannot be inferred.
+    #[salsa::tracked]
+    pub(crate) fn try_metaclass(self, db: &'db dyn Db) -> Result<Type<'db>, MetaclassError<'db>> {
+        /// Infer the metaclass of a class, tracking the classes that have been visited to detect
+        /// cyclic definitions.
+        fn infer<'db>(
+            db: &'db dyn Db,
+            class: Class<'db>,
+            seen: &mut SeenSet<Class<'db>>,
+        ) -> Result<Type<'db>, MetaclassError<'db>> {
+            // Recursively infer the metaclass of a class, ensuring that cyclic definitions are
+            // detected.
+            let mut safe_recurse = |class: Class<'db>| -> Result<Type<'db>, MetaclassError<'db>> {
+                // Each base must be considered in isolation.
+                let num_seen = seen.len();
+                if !seen.insert(class) {
+                    return Err(MetaclassError {
+                        kind: MetaclassErrorKind::CyclicDefinition,
+                    });
+                }
+                let metaclass = infer(db, class, seen)?;
+                seen.truncate(num_seen);
+                Ok(metaclass)
+            };
+
+            let mut base_classes = class
+                .explicit_bases(db)
+                .iter()
+                .copied()
+                .filter_map(Type::into_class_literal);
+
+            // Identify the class's own metaclass (or take the first base class's metaclass).
+            let metaclass = if let Some(metaclass) = class.explicit_metaclass(db) {
+                metaclass
+            } else if let Some(base_class) = base_classes.next() {
+                safe_recurse(base_class.class)?
+            } else {
+                KnownClass::Type.to_class(db)
+            };
+
+            let Type::ClassLiteral(mut candidate) = metaclass else {
+                // If the metaclass is not a class, return it directly.
+                return Ok(metaclass);
+            };
+
+            // Reconcile all base classes' metaclasses with the candidate metaclass.
+            //
+            // See:
+            // - https://docs.python.org/3/reference/datamodel.html#determining-the-appropriate-metaclass
+            // - https://github.com/python/cpython/blob/83ba8c2bba834c0b92de669cac16fcda17485e0e/Objects/typeobject.c#L3629-L3663
+            for base_class in base_classes {
+                let metaclass = safe_recurse(base_class.class)?;
+                let Type::ClassLiteral(metaclass) = metaclass else {
+                    continue;
+                };
+                if metaclass.class.is_subclass_of(db, candidate.class) {
+                    candidate = metaclass;
+                    continue;
+                }
+                if candidate.class.is_subclass_of(db, metaclass.class) {
+                    continue;
+                }
+                return Err(MetaclassError {
+                    kind: MetaclassErrorKind::Conflict {
+                        metaclass1: candidate.class,
+                        metaclass2: metaclass.class,
+                    },
+                });
+            }
+
+            Ok(Type::ClassLiteral(candidate))
+        }
+
+        infer(db, self, &mut SeenSet::new(self))
+    }
+
     /// Returns the class member of this class named `name`.
     ///
     /// The member resolves to a member on the class itself or any of its proper superclasses.
@@ -2038,6 +2147,10 @@ impl<'db> Class<'db> {
             );
         }
 
+        if name == "__class__" {
+            return self.metaclass(db).into();
+        }
+
         for superclass in self.iter_mro(db) {
             match superclass {
                 // TODO we may instead want to record the fact that we encountered dynamic, and intersect it with
@@ -2068,6 +2181,38 @@ impl<'db> Class<'db> {
     }
 }
 
+/// A utility struct for detecting duplicates in class hierarchies while storing the initial
+/// entry on the stack.
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct SeenSet<T: Hash + Eq> {
+    initial: T,
+    visited: IndexSet<T>,
+}
+
+impl<T: Hash + Eq> SeenSet<T> {
+    fn new(initial: T) -> SeenSet<T> {
+        Self {
+            initial,
+            visited: IndexSet::new(),
+        }
+    }
+
+    fn len(&self) -> usize {
+        self.visited.len()
+    }
+
+    fn truncate(&mut self, len: usize) {
+        self.visited.truncate(len);
+    }
+
+    fn insert(&mut self, value: T) -> bool {
+        if value == self.initial {
+            return false;
+        }
+        self.visited.insert(value)
+    }
+}
+
 /// A singleton type representing a single class object at runtime.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct ClassLiteralType<'db> {
@@ -2130,6 +2275,36 @@ impl<'db> From<InstanceType<'db>> for Type<'db> {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(super) struct MetaclassError<'db> {
+    kind: MetaclassErrorKind<'db>,
+}
+
+impl<'db> MetaclassError<'db> {
+    /// Return an [`MetaclassErrorKind`] variant describing why we could not resolve the metaclass for this class.
+    pub(super) fn reason(&self) -> &MetaclassErrorKind<'db> {
+        &self.kind
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(super) enum MetaclassErrorKind<'db> {
+    /// The class has incompatible metaclasses in its inheritance hierarchy.
+    ///
+    /// The metaclass of a derived class must be a (non-strict) subclass of the metaclasses of all
+    /// its bases.
+    Conflict {
+        metaclass1: Class<'db>,
+        metaclass2: Class<'db>,
+    },
+    /// The class inherits from itself!
+    ///
+    /// This is very unlikely to happen in working real-world code,
+    /// but it's important to explicitly account for it.
+    /// If we don't, there's a possibility of an infinite loop and a panic.
+    CyclicDefinition,
+}
+
 #[salsa::interned]
 pub struct UnionType<'db> {
     /// The union type includes values in any of these types.

crates/red_knot_python_semantic/src/types/infer.rs

@@ -57,9 +57,9 @@ use crate::types::unpacker::{UnpackResult, Unpacker};
 use crate::types::{
     bindings_ty, builtins_symbol, declarations_ty, global_symbol, symbol, typing_extensions_symbol,
     Boundness, BytesLiteralType, Class, ClassLiteralType, FunctionType, InstanceType,
-    IterationOutcome, KnownClass, KnownFunction, KnownInstance, SliceLiteralType,
+    IterationOutcome, KnownClass, KnownFunction, KnownInstance, MetaclassErrorKind,
-    StringLiteralType, Symbol, Truthiness, TupleType, Type, TypeArrayDisplay, UnionBuilder,
+    SliceLiteralType, StringLiteralType, Symbol, Truthiness, TupleType, Type, TypeArrayDisplay,
-    UnionType,
+    UnionBuilder, UnionType,
 };
 use crate::unpack::Unpack;
 use crate::util::subscript::{PyIndex, PySlice};
@@ -450,9 +450,11 @@ impl<'db> TypeInferenceBuilder<'db> {
     }
 
     /// Iterate over all class definitions to check that Python will be able to create a
-    /// consistent "[method resolution order]" for each class at runtime. If not, issue a diagnostic.
+    /// consistent "[method resolution order]" and [metaclass] for each class at runtime. If not,
+    /// issue a diagnostic.
     ///
     /// [method resolution order]: https://docs.python.org/3/glossary.html#term-method-resolution-order
+    /// [metaclass]: https://docs.python.org/3/reference/datamodel.html#metaclasses
     fn check_class_definitions(&mut self) {
         let class_definitions = self
             .types
@@ -461,15 +463,8 @@ impl<'db> TypeInferenceBuilder<'db> {
             .filter_map(|ty| ty.into_class_literal())
             .map(|class_ty| class_ty.class);
 
-        let invalid_mros = class_definitions.filter_map(|class| {
+        for class in class_definitions {
-            class
+            if let Err(mro_error) = class.try_mro(self.db).as_ref() {
-                .try_mro(self.db)
-                .as_ref()
-                .err()
-                .map(|mro_error| (class, mro_error))
-        });
-
-        for (class, mro_error) in invalid_mros {
                 match mro_error.reason() {
                     MroErrorKind::DuplicateBases(duplicates) => {
                         let base_nodes = class.node(self.db).bases();
@@ -477,7 +472,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                             self.diagnostics.add(
                                 (&base_nodes[*index]).into(),
                                 "duplicate-base",
-                             format_args!("Duplicate base class `{}`", duplicate.name(self.db))
+                                format_args!("Duplicate base class `{}`", duplicate.name(self.db)),
                             );
                         }
                     }
@@ -488,7 +483,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                             "Cyclic definition of `{}` or bases of `{}` (class cannot inherit from itself)",
                             class.name(self.db),
                             class.name(self.db)
-                    )
+                        ),
                     ),
                     MroErrorKind::InvalidBases(bases) => {
                         let base_nodes = class.node(self.db).bases();
@@ -499,19 +494,42 @@ impl<'db> TypeInferenceBuilder<'db> {
                                 format_args!(
                                     "Invalid class base with type `{}` (all bases must be a class, `Any`, `Unknown` or `Todo`)",
                                     base_ty.display(self.db)
-                            )
+                                ),
                             );
                         }
-                },
+                    }
-                MroErrorKind::UnresolvableMro{bases_list} => self.diagnostics.add(
+                    MroErrorKind::UnresolvableMro { bases_list } => self.diagnostics.add(
                         class.node(self.db).into(),
                         "inconsistent-mro",
                         format_args!(
                             "Cannot create a consistent method resolution order (MRO) for class `{}` with bases list `[{}]`",
                             class.name(self.db),
                             bases_list.iter().map(|base| base.display(self.db)).join(", ")
+                        ),
                     )
-                )
+                }
+            }
+
+            if let Err(metaclass_error) = class.try_metaclass(self.db) {
+                match metaclass_error.reason() {
+                    MetaclassErrorKind::Conflict {
+                        metaclass1,
+                        metaclass2
+                    } => self.diagnostics.add(
+                        class.node(self.db).into(),
+                        "conflicting-metaclass",
+                        format_args!(
+                            "The metaclass of a derived class (`{}`) must be a subclass of the metaclasses of all its bases, but `{}` and `{}` have no subclass relationship",
+                            class.name(self.db),
+                            metaclass1.name(self.db),
+                            metaclass2.name(self.db),
+                        ),
+                    ),
+                    MetaclassErrorKind::CyclicDefinition => {
+                        // Cyclic class definition diagnostic will already have been emitted above
+                        // in MRO calculation.
+                    }
+                }
             }
         }
     }
@@ -1187,9 +1205,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                             context_expression.into(),
                             "invalid-context-manager",
                             format_args!("
-                                Object of type `{context_expression}` cannot be used with `with` because the method `__enter__` of type `{enter_ty}` is not callable",
+                                Object of type `{context_expression}` cannot be used with `with` because the method `__enter__` of type `{enter_ty}` is not callable", context_expression = context_expression_ty.display(self.db), enter_ty = enter_ty.display(self.db)
-                                context_expression = context_expression_ty.display(self.db),
-                                enter_ty = enter_ty.display(self.db)
                             ),
                         );
                         err.return_ty()
@@ -4403,7 +4419,6 @@ fn perform_membership_test_comparison<'db>(
 
 #[cfg(test)]
 mod tests {
-
     use anyhow::Context;
 
     use crate::db::tests::TestDb;