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,6 +4637,14 @@ impl<'db> TypeInferenceBuilder<'db> {
             }
         }
 
+        let class = match callable_type {
+            Type::ClassLiteral(class) => Some(ClassType::NonGeneric(class)),
+            Type::GenericAlias(generic) => Some(ClassType::Generic(generic)),
+            Type::SubclassOf(subclass) => subclass.subclass_of().into_class(),
+            _ => None,
+        };
+
+        if let Some(class) = class {
             // 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
             // can be called even though class-literals cannot. This is because even though a protocol class
@@ -4644,16 +4652,18 @@ impl<'db> TypeInferenceBuilder<'db> {
             // 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 {
+            if !callable_type.is_subclass_of() {
-            Type::ClassLiteral(class) => Some(class),
+                if let Some(protocol) = class
-            Type::GenericAlias(generic) => Some(generic.origin(self.db())),
+                    .class_literal(self.db())
-            _ => None,
+                    .0
-        };
+                    .into_protocol_class(self.db())
-
-        if let Some(protocol) =
-            possible_protocol_class.and_then(|class| class.into_protocol_class(self.db()))
                 {
-            report_attempted_protocol_instantiation(&self.context, call_expression, protocol);
+                    report_attempted_protocol_instantiation(
+                        &self.context,
+                        call_expression,
+                        protocol,
+                    );
+                }
             }
 
             // For class literals we model the entire class instantiation logic, so it is handled
@@ -4661,22 +4671,8 @@ impl<'db> TypeInferenceBuilder<'db> {
             // the `try_call` path below.
             // 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.
-        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,6 +4683,12 @@ impl<'db> TypeInferenceBuilder<'db> {
                         | KnownClass::TypeVar
                 )
             )
+            // temporary special-casing for all subclasses of `enum.Enum`
+            // until we support the functional syntax for creating enum classes
+            && KnownClass::Enum
+                .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 =
@@ -4699,6 +4701,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                         err.return_type()
                     });
             }
+        }
 
         let signatures = callable_type.signatures(self.db());
         let bindings = Bindings::match_parameters(signatures, &call_arguments);