[red-knot] More precise inference for classes with non-class metaclasses

crates/red_knot_python_semantic/resources/mdtest/metaclass.md

@@ -170,8 +170,35 @@ def f(*args, **kwargs) -> int: ...
 
 class A(metaclass=f): ...
 
-# TODO should be `type[int]`
-reveal_type(A.__class__)  # revealed: @Todo(metaclass not a class)
+# TODO: Should be `int`
+reveal_type(A)  # revealed: Literal[A]
+reveal_type(A.__class__)  # revealed: type[int]
+
+def _(n: int):
+    # error: [invalid-metaclass]
+    class B(metaclass=n): ...
+    # TODO: Should be `Unknown`
+    reveal_type(B)  # revealed: Literal[B]
+    reveal_type(B.__class__)  # revealed: type[Unknown]
+
+def _(flag: bool):
+    m = f if flag else 42
+
+    # error: [invalid-metaclass]
+    class C(metaclass=m): ...
+    # TODO: Should be `int | Unknown`
+    reveal_type(C)  # revealed: Literal[C]
+    reveal_type(C.__class__)  # revealed: type[Unknown]
+
+class SignatureMismatch: ...
+
+# TODO: Emit a diagnostic
+class D(metaclass=SignatureMismatch): ...
+
+# TODO: Should be `Unknown`
+reveal_type(D)  # revealed: Literal[D]
+# TODO: Should be `type[Unknown]`
+reveal_type(D.__class__)  # revealed: Literal[SignatureMismatch]
 ```
 
 ## Cyclic

crates/red_knot_python_semantic/src/types.rs

@@ -3607,10 +3607,60 @@ impl<'db> Class<'db> {
                 explicit_metaclass_of: class_metaclass_was_from,
             }
         } else {
-            // TODO: If the metaclass is not a class, we should verify that it's a callable
-            // which accepts the same arguments as `type.__new__` (otherwise error), and return
-            // the meta-type of its return type. (And validate that is a class type?)
-            return Ok(todo_type!("metaclass not a class"));
+            let self_ty = Type::class_literal(self);
+            let bases = TupleType::from_elements(db, self.explicit_bases(db));
+            // TODO: Should be `dict[str, Any]`
+            let namespace = KnownClass::Dict.to_instance(db);
+
+            // TODO: Other keyword arguments?
+            let arguments = CallArguments::positional([self_ty, bases, namespace]);
+
+            let return_ty_result = match metaclass.call(db, &arguments) {
+                CallOutcome::NotCallable { not_callable_ty } => Err(MetaclassError {
+                    kind: MetaclassErrorKind::NotCallable(not_callable_ty),
+                }),
+
+                CallOutcome::Union {
+                    outcomes,
+                    called_ty,
+                } => {
+                    let mut partly_not_callable = false;
+
+                    let return_ty = outcomes
+                        .iter()
+                        .fold(None, |acc, outcome| {
+                            let ty = outcome.return_ty(db);
+
+                            match (acc, ty) {
+                                (acc, None) => {
+                                    partly_not_callable = true;
+                                    acc
+                                }
+                                (None, Some(ty)) => Some(UnionBuilder::new(db).add(ty)),
+                                (Some(builder), Some(ty)) => Some(builder.add(ty)),
+                            }
+                        })
+                        .map(UnionBuilder::build);
+
+                    if partly_not_callable {
+                        Err(MetaclassError {
+                            kind: MetaclassErrorKind::PartlyNotCallable(called_ty),
+                        })
+                    } else {
+                        Ok(return_ty.unwrap_or(Type::Unknown))
+                    }
+                }
+
+                CallOutcome::PossiblyUnboundDunderCall { called_ty, .. } => Err(MetaclassError {
+                    kind: MetaclassErrorKind::PartlyNotCallable(called_ty),
+                }),
+
+                CallOutcome::Callable { binding }
+                | CallOutcome::RevealType { binding, .. }
+                | CallOutcome::StaticAssertionError { binding, .. } => Ok(binding.return_ty()),
+            };
+
+            return return_ty_result.map(|ty| ty.to_meta_type(db));
         };
 
         // Reconcile all base classes' metaclasses with the candidate metaclass.
@@ -3815,6 +3865,10 @@ pub(super) enum MetaclassErrorKind<'db> {
         /// inferred metaclass of a base class. This helps us give better error messages in diagnostics.
         candidate1_is_base_class: bool,
     },
+    /// The metaclass is not callable
+    NotCallable(Type<'db>),
+    /// The metaclass is of a union type whose some members are not callable
+    PartlyNotCallable(Type<'db>),
 }
 
 #[salsa::interned]

crates/red_knot_python_semantic/src/types/diagnostic.rs

@@ -36,6 +36,7 @@ pub(crate) fn register_lints(registry: &mut LintRegistryBuilder) {
     registry.register_lint(&INVALID_CONTEXT_MANAGER);
     registry.register_lint(&INVALID_DECLARATION);
     registry.register_lint(&INVALID_EXCEPTION_CAUGHT);
+    registry.register_lint(&INVALID_METACLASS);
     registry.register_lint(&INVALID_PARAMETER_DEFAULT);
     registry.register_lint(&INVALID_RAISE);
     registry.register_lint(&INVALID_TYPE_FORM);
@@ -290,6 +291,7 @@ declare_lint! {
 }
 
 declare_lint! {
+    /// ## What it does
     /// Checks for exception handlers that catch non-exception classes.
     ///
     /// ## Why is this bad?
@@ -324,6 +326,33 @@ declare_lint! {
     }
 }
 
+declare_lint! {
+    /// ## What it does
+    /// Checks for arguments to `metaclass=` that are invalid.
+    ///
+    /// ## Why is this bad?
+    /// Python allows arbitrary expressions to be used as the argument to `metaclass=`.
+    /// These expressions, however, need to be callable and accept the same arguments
+    /// as `type.__new__`.
+    ///
+    /// ## Example
+    ///
+    /// ```python
+    /// def f(): ...
+    ///
+    /// # TypeError: f() takes 0 positional arguments but 3 were given
+    /// class B(metaclass=f): ...
+    /// ```
+    ///
+    /// ## References
+    /// - [Python documentation: Metaclasses](https://docs.python.org/3/reference/datamodel.html#metaclasses)
+    pub(crate) static INVALID_METACLASS = {
+        summary: "detects invalid `metaclass=` arguments",
+        status: LintStatus::preview("1.0.0"),
+        default_level: Level::Error,
+    }
+}
+
 declare_lint! {
     /// ## What it does
     /// Checks for default values that can't be assigned to the parameter's annotated type.

crates/red_knot_python_semantic/src/types/infer.rs

@@ -77,7 +77,7 @@ use super::diagnostic::{
     report_index_out_of_bounds, report_invalid_exception_caught, report_invalid_exception_cause,
     report_invalid_exception_raised, report_non_subscriptable,
     report_possibly_unresolved_reference, report_slice_step_size_zero, report_unresolved_reference,
-    SUBCLASS_OF_FINAL_CLASS,
+    INVALID_METACLASS, SUBCLASS_OF_FINAL_CLASS,
 };
 use super::slots::check_class_slots;
 use super::string_annotation::{
@@ -639,6 +639,19 @@ impl<'db> TypeInferenceBuilder<'db> {
             // (4) Check that the class's metaclass can be determined without error.
             if let Err(metaclass_error) = class.try_metaclass(self.db()) {
                 match metaclass_error.reason() {
+                    MetaclassErrorKind::NotCallable(ty) => self.context.report_lint(
+                        &INVALID_METACLASS,
+                        class_node.into(),
+                        format_args!("Metaclass type `{}` is not callable", ty.display(self.db())),
+                    ),
+                    MetaclassErrorKind::PartlyNotCallable(ty) => self.context.report_lint(
+                        &INVALID_METACLASS,
+                        class_node.into(),
+                        format_args!(
+                            "Metaclass type `{}` is partly not callable",
+                            ty.display(self.db())
+                        ),
+                    ),
                     MetaclassErrorKind::Conflict {
                         candidate1:
                             MetaclassCandidate {