[ty] handle more enum.auto() cases

Author: thejchap

crates/ty_python_semantic/resources/mdtest/enums.md

@@ -320,6 +320,11 @@ reveal_type(enum_members(Answer))
 
 reveal_type(Answer.YES.value)  # revealed: Literal[1]
 reveal_type(Answer.NO.value)  # revealed: Literal[2]
+
+class SingleMember(Enum):
+    SINGLE = auto()
+
+reveal_type(SingleMember.SINGLE.value)  # revealed: Literal[1]
 ```
 
 Usages of `auto()` can be combined with manual value assignments:
@@ -348,6 +353,11 @@ class Answer(StrEnum):
 
 reveal_type(Answer.YES.value)  # revealed: Literal["yes"]
 reveal_type(Answer.NO.value)  # revealed: Literal["no"]
+
+class SingleMember(StrEnum):
+    SINGLE = auto()
+
+reveal_type(SingleMember.SINGLE.value)  # revealed: Literal["single"]
 ```
 
 Using `auto()` with `IntEnum` also works as expected:
@@ -363,6 +373,36 @@ reveal_type(Answer.YES.value)  # revealed: Literal[1]
 reveal_type(Answer.NO.value)  # revealed: Literal[2]
 ```
 
+Using `auto()` with non-integer mixins:
+
+```python
+from enum import Enum, auto
+
+class A(str, Enum):
+    X = auto()
+    Y = auto()
+
+reveal_type(A.X.value)  # revealed: Any
+
+class B(bytes, Enum):
+    X = auto()
+    Y = auto()
+
+reveal_type(B.X.value)  # revealed: Any
+
+class C(tuple, Enum):
+    X = auto()
+    Y = auto()
+
+reveal_type(C.X.value)  # revealed: Any
+
+class D(float, Enum):
+    X = auto()
+    Y = auto()
+
+reveal_type(D.X.value)  # revealed: Any
+```
+
 Combining aliases with `auto()`:
 
 ```py

crates/ty_python_semantic/src/types.rs

@@ -4365,6 +4365,16 @@ impl<'db> Type<'db> {
                 Place::bound(todo_type!("ParamSpecArgs / ParamSpecKwargs")).into()
             }
 
+            Type::NominalInstance(instance)
+                if matches!(name_str, "value" | "_value_")
+                    && is_single_member_enum(db, instance.class(db).class_literal(db).0) =>
+            {
+                enum_metadata(db, instance.class(db).class_literal(db).0)
+                    .and_then(|metadata| metadata.members.get_index(0).map(|(_, v)| *v))
+                    .map_or(Place::Undefined, Place::bound)
+                    .into()
+            }
+
             Type::NominalInstance(..)
             | Type::ProtocolInstance(..)
             | Type::BooleanLiteral(..)

crates/ty_python_semantic/src/types/class.rs

@@ -3778,6 +3778,7 @@ pub enum KnownClass {
     Member,
     Nonmember,
     StrEnum,
+    IntEnum,
     // abc
     ABCMeta,
     // Types
@@ -3925,6 +3926,7 @@ impl KnownClass {
             | Self::Member
             | Self::Nonmember
             | Self::StrEnum
+            | Self::IntEnum
             | Self::ABCMeta
             | Self::Iterable
             | Self::Iterator
@@ -3986,6 +3988,7 @@ impl KnownClass {
             | KnownClass::Member
             | KnownClass::Nonmember
             | KnownClass::StrEnum
+            | KnownClass::IntEnum
             | KnownClass::ABCMeta
             | KnownClass::GenericAlias
             | KnownClass::ModuleType
@@ -4069,6 +4072,7 @@ impl KnownClass {
             | KnownClass::Member
             | KnownClass::Nonmember
             | KnownClass::StrEnum
+            | KnownClass::IntEnum
             | KnownClass::ABCMeta
             | KnownClass::GenericAlias
             | KnownClass::ModuleType
@@ -4152,6 +4156,7 @@ impl KnownClass {
             | KnownClass::Member
             | KnownClass::Nonmember
             | KnownClass::StrEnum
+            | KnownClass::IntEnum
             | KnownClass::ABCMeta
             | KnownClass::GenericAlias
             | KnownClass::ModuleType
@@ -4274,6 +4279,7 @@ impl KnownClass {
             | Self::Member
             | Self::Nonmember
             | Self::StrEnum
+            | Self::IntEnum
             | Self::ABCMeta
             | Self::Super
             | Self::StdlibAlias
@@ -4330,6 +4336,7 @@ impl KnownClass {
             | KnownClass::Member
             | KnownClass::Nonmember
             | KnownClass::StrEnum
+            | KnownClass::IntEnum
             | KnownClass::ABCMeta
             | KnownClass::GenericAlias
             | KnownClass::ModuleType
@@ -4447,6 +4454,7 @@ impl KnownClass {
             Self::Member => "member",
             Self::Nonmember => "nonmember",
             Self::StrEnum => "StrEnum",
+            Self::IntEnum => "IntEnum",
             Self::ABCMeta => "ABCMeta",
             Self::Super => "super",
             Self::Iterable => "Iterable",
@@ -4719,7 +4727,8 @@ impl KnownClass {
             | Self::Auto
             | Self::Member
             | Self::Nonmember
-            | Self::StrEnum => KnownModule::Enum,
+            | Self::StrEnum
+            | Self::IntEnum => KnownModule::Enum,
             Self::GenericAlias
             | Self::ModuleType
             | Self::FunctionType
@@ -4850,6 +4859,7 @@ impl KnownClass {
             | Self::Member
             | Self::Nonmember
             | Self::StrEnum
+            | Self::IntEnum
             | Self::ABCMeta
             | Self::Super
             | Self::NewType
@@ -4937,6 +4947,7 @@ impl KnownClass {
             | Self::Member
             | Self::Nonmember
             | Self::StrEnum
+            | Self::IntEnum
             | Self::ABCMeta
             | Self::Super
             | Self::UnionType
@@ -5106,6 +5117,7 @@ impl KnownClass {
             | Self::Member
             | Self::Nonmember
             | Self::StrEnum
+            | Self::IntEnum
             | Self::ABCMeta
             | Self::Super
             | Self::NotImplementedType

crates/ty_python_semantic/src/types/enums.rs

@@ -68,9 +68,6 @@ pub(crate) fn enum_metadata<'db>(
         return None;
     }
 
-    let is_str_enum =
-        Type::ClassLiteral(class).is_subtype_of(db, KnownClass::StrEnum.to_subclass_of(db));
-
     let scope_id = class.body_scope(db);
     let use_def_map = use_def_map(db, scope_id);
     let table = place_table(db, scope_id);
@@ -141,14 +138,38 @@ pub(crate) fn enum_metadata<'db>(
                             // enum.auto
                             Some(KnownClass::Auto) => {
                                 auto_counter += 1;
-                                Some(if is_str_enum {
+                                let auto_value_ty = if Type::ClassLiteral(class)
+                                    .is_subtype_of(db, KnownClass::StrEnum.to_subclass_of(db))
+                                {
                                     Type::StringLiteral(StringLiteralType::new(
                                         db,
                                         name.to_lowercase().as_str(),
                                     ))
-                                } else {
+                                } else if Type::ClassLiteral(class)
+                                    .is_subtype_of(db, KnownClass::IntEnum.to_subclass_of(db))
+                                {
                                     Type::IntLiteral(auto_counter)
-                                })
+                                } else {
+                                    let has_custom_mixin =
+                                        class.iter_mro(db, None).skip(1).any(|base| {
+                                            base.into_class().is_some_and(|class| {
+                                                !(class.is_object(db)
+                                                    || KnownClass::Enum
+                                                        .to_subclass_of(db)
+                                                        .to_class_type(db)
+                                                        .is_some_and(|enum_class| {
+                                                            class.is_subclass_of(db, enum_class)
+                                                        })
+                                                    || class.is_known(db, KnownClass::Enum))
+                                            })
+                                        });
+                                    if has_custom_mixin {
+                                        Type::any()
+                                    } else {
+                                        Type::IntLiteral(auto_counter)
+                                    }
+                                };
+                                Some(auto_value_ty)
                             }
 
                             _ => None,