Generalize special-casing for enums constructed with the functional syntax (#17885)

crates/ty_python_semantic/resources/mdtest/annotations/unsupported_special_types.md

@@ -8,9 +8,8 @@ import collections
 import enum
 import typing
 
-# TODO: should not error (requires understanding metaclass `__call__`)
+MyEnum = enum.Enum("MyEnum", ["foo", "bar", "baz"])
-MyEnum = enum.Enum("MyEnum", ["foo", "bar", "baz"])  # error: [too-many-positional-arguments]
+MyIntEnum = enum.IntEnum("MyIntEnum", ["foo", "bar", "baz"])
-
 MyTypedDict = typing.TypedDict("MyTypedDict", {"foo": int})
 MyNamedTuple1 = typing.NamedTuple("MyNamedTuple1", [("foo", int)])
 MyNamedTuple2 = collections.namedtuple("MyNamedTuple2", ["foo"])

crates/ty_python_semantic/src/types.rs

@@ -3890,6 +3890,9 @@ impl<'db> Type<'db> {
                     );
                     Signatures::single(signature)
                 }
+                Some(KnownClass::Enum) => {
+                    Signatures::single(CallableSignature::todo("functional `Enum` syntax"))
+                }
 
                 Some(KnownClass::Super) => {
                     // ```py
@@ -4836,14 +4839,6 @@ impl<'db> Type<'db> {
                 Some(KnownClass::NamedTuple) => Ok(todo_type!(
                     "Support for functional `typing.NamedTuple` syntax"
                 )),
-                _ if instance
-                    .class()
-                    .iter_mro(db)
-                    .filter_map(ClassBase::into_class)
-                    .any(|class| class.is_known(db, KnownClass::Enum)) =>
-                {
-                    Ok(todo_type!("Support for functional `enum` syntax"))
-                }
                 _ => Err(InvalidTypeExpressionError {
                     invalid_expressions: smallvec::smallvec![InvalidTypeExpression::InvalidType(
                         *self

crates/ty_python_semantic/src/types/infer.rs

@@ -4637,46 +4637,42 @@ impl<'db> TypeInferenceBuilder<'db> {
             }
         }
 
-        // It might look odd here that we emit an error for class-literals and generic aliases but not
+        let class = match callable_type {
-        // `type[]` types. But it's deliberate! The typing spec explicitly mandates that `type[]` types
+            Type::ClassLiteral(class) => Some(ClassType::NonGeneric(class)),
-        // can be called even though class-literals cannot. This is because even though a protocol class
+            Type::GenericAlias(generic) => Some(ClassType::Generic(generic)),
-        // `SomeProtocol` is always an abstract class, `type[SomeProtocol]` can be a concrete subclass of
+            Type::SubclassOf(subclass) => subclass.subclass_of().into_class(),
-        // that protocol -- and indeed, according to the spec, type checkers must disallow abstract
-        // subclasses of the protocol to be passed to parameters that accept `type[SomeProtocol]`.
-        // <https://typing.python.org/en/latest/spec/protocol.html#type-and-class-objects-vs-protocols>.
-        let possible_protocol_class = match callable_type {
-            Type::ClassLiteral(class) => Some(class),
-            Type::GenericAlias(generic) => Some(generic.origin(self.db())),
             _ => None,
         };
 
-        if let Some(protocol) =
+        if let Some(class) = class {
-            possible_protocol_class.and_then(|class| class.into_protocol_class(self.db()))
+            // It might look odd here that we emit an error for class-literals and generic aliases but not
-        {
+            // `type[]` types. But it's deliberate! The typing spec explicitly mandates that `type[]` types
-            report_attempted_protocol_instantiation(&self.context, call_expression, protocol);
+            // can be called even though class-literals cannot. This is because even though a protocol class
-        }
+            // `SomeProtocol` is always an abstract class, `type[SomeProtocol]` can be a concrete subclass of
+            // that protocol -- and indeed, according to the spec, type checkers must disallow abstract
+            // subclasses of the protocol to be passed to parameters that accept `type[SomeProtocol]`.
+            // <https://typing.python.org/en/latest/spec/protocol.html#type-and-class-objects-vs-protocols>.
+            if !callable_type.is_subclass_of() {
+                if let Some(protocol) = class
+                    .class_literal(self.db())
+                    .0
+                    .into_protocol_class(self.db())
+                {
+                    report_attempted_protocol_instantiation(
+                        &self.context,
+                        call_expression,
+                        protocol,
+                    );
+                }
+            }
 
-        // For class literals we model the entire class instantiation logic, so it is handled
+            // For class literals we model the entire class instantiation logic, so it is handled
-        // in a separate function. For some known classes we have manual signatures defined and use
+            // in a separate function. For some known classes we have manual signatures defined and use
-        // the `try_call` path below.
+            // the `try_call` path below.
-        // TODO: it should be possible to move these special cases into the `try_call_constructor`
+            // TODO: it should be possible to move these special cases into the `try_call_constructor`
-        // path instead, or even remove some entirely once we support overloads fully.
+            // path instead, or even remove some entirely once we support overloads fully.
-        let (call_constructor, known_class) = match callable_type {
+            if !matches!(
-            Type::ClassLiteral(class) => (true, class.known(self.db())),
+                class.known(self.db()),
-            Type::GenericAlias(generic) => (true, ClassType::Generic(generic).known(self.db())),
-            Type::SubclassOf(subclass) => (
-                true,
-                subclass
-                    .subclass_of()
-                    .into_class()
-                    .and_then(|class| class.known(self.db())),
-            ),
-            _ => (false, None),
-        };
-
-        if call_constructor
-            && !matches!(
-                known_class,
                 Some(
                     KnownClass::Bool
                         | KnownClass::Str
@@ -4687,17 +4683,24 @@ impl<'db> TypeInferenceBuilder<'db> {
                         | KnownClass::TypeVar
                 )
             )
-        {
+            // temporary special-casing for all subclasses of `enum.Enum`
-            let argument_forms = vec![Some(ParameterForm::Value); call_arguments.len()];
+            // until we support the functional syntax for creating enum classes
-            let call_argument_types =
+            && KnownClass::Enum
-                self.infer_argument_types(arguments, call_arguments, &argument_forms);
+                .to_class_literal(self.db())
+                .to_class_type(self.db())
+                .is_none_or(|enum_class| !class.is_subclass_of(self.db(), enum_class))
+            {
+                let argument_forms = vec![Some(ParameterForm::Value); call_arguments.len()];
+                let call_argument_types =
+                    self.infer_argument_types(arguments, call_arguments, &argument_forms);
 
-            return callable_type
+                return callable_type
-                .try_call_constructor(self.db(), call_argument_types)
+                    .try_call_constructor(self.db(), call_argument_types)
-                .unwrap_or_else(|err| {
+                    .unwrap_or_else(|err| {
-                    err.report_diagnostic(&self.context, callable_type, call_expression.into());
+                        err.report_diagnostic(&self.context, callable_type, call_expression.into());
-                    err.return_type()
+                        err.return_type()
-                });
+                    });
+            }
         }
 
         let signatures = callable_type.signatures(self.db());