[ty] tuple syntax and fix flag/intenum/etc routing

astral-sh/ty#876

Author: thejchap

crates/ty_python_semantic/resources/mdtest/enums.md

@@ -1288,13 +1288,13 @@ from ty_extensions import enum_members
 
 Color = Enum("Color", [("RED", 1), ("GREEN", 2), ("BLUE", 3)])
 
-# TODO: this should be supported
-reveal_type(enum_members(Color))  # revealed: Unknown
+# revealed: tuple[Literal["RED"], Literal["GREEN"], Literal["BLUE"]]
+reveal_type(enum_members(Color))
 
 Color = Enum("Color", (("RED", 1), ("GREEN", 2), ("BLUE", 3)))
 
-# TODO: this should be supported
-reveal_type(enum_members(Color))  # revealed: Unknown
+# revealed: tuple[Literal["RED"], Literal["GREEN"], Literal["BLUE"]]
+reveal_type(enum_members(Color))
 ```
 
 ### List of strings
@@ -1387,6 +1387,42 @@ reveal_type(Http.OK.value)  # revealed: Literal[1]
 reveal_type(Http.NOT_FOUND.value)  # revealed: Literal[2]
 ```
 
+### IntEnum function syntax
+
+```py
+from enum import IntEnum
+from ty_extensions import enum_members
+
+Color = IntEnum("Color", "RED GREEN BLUE")
+
+# revealed: tuple[Literal["RED"], Literal["GREEN"], Literal["BLUE"]]
+reveal_type(enum_members(Color))
+```
+
+### Flag function syntax
+
+```py
+from enum import Flag
+from ty_extensions import enum_members
+
+Perm = Flag("Perm", "READ WRITE EXECUTE")
+
+# revealed: tuple[Literal["READ"], Literal["WRITE"], Literal["EXECUTE"]]
+reveal_type(enum_members(Perm))
+```
+
+### IntFlag function syntax
+
+```py
+from enum import IntFlag
+from ty_extensions import enum_members
+
+Perm = IntFlag("Perm", "READ WRITE EXECUTE")
+
+# revealed: tuple[Literal["READ"], Literal["WRITE"], Literal["EXECUTE"]]
+reveal_type(enum_members(Perm))
+```
+
 ## Exhaustiveness checking
 
 ## `if` statements

crates/ty_python_semantic/src/types/class/known.rs

@@ -68,6 +68,9 @@ pub enum KnownClass {
     Member,
     Nonmember,
     StrEnum,
+    IntEnum,
+    Flag,
+    IntFlag,
     // abc
     ABCMeta,
     // Types
@@ -225,6 +228,9 @@ impl KnownClass {
             | Self::Member
             | Self::Nonmember
             | Self::StrEnum
+            | Self::IntEnum
+            | Self::Flag
+            | Self::IntFlag
             | Self::ABCMeta
             | Self::Iterable
             | Self::Iterator
@@ -294,6 +300,9 @@ impl KnownClass {
             | KnownClass::Member
             | KnownClass::Nonmember
             | KnownClass::StrEnum
+            | KnownClass::IntEnum
+            | KnownClass::Flag
+            | KnownClass::IntFlag
             | KnownClass::ABCMeta
             | KnownClass::GenericAlias
             | KnownClass::ModuleType
@@ -385,6 +394,9 @@ impl KnownClass {
             | KnownClass::Member
             | KnownClass::Nonmember
             | KnownClass::StrEnum
+            | KnownClass::IntEnum
+            | KnownClass::Flag
+            | KnownClass::IntFlag
             | KnownClass::ABCMeta
             | KnownClass::GenericAlias
             | KnownClass::ModuleType
@@ -476,6 +488,9 @@ impl KnownClass {
             | KnownClass::Member
             | KnownClass::Nonmember
             | KnownClass::StrEnum
+            | KnownClass::IntEnum
+            | KnownClass::Flag
+            | KnownClass::IntFlag
             | KnownClass::ABCMeta
             | KnownClass::GenericAlias
             | KnownClass::ModuleType
@@ -608,6 +623,9 @@ impl KnownClass {
             | Self::Member
             | Self::Nonmember
             | Self::StrEnum
+            | Self::IntEnum
+            | Self::Flag
+            | Self::IntFlag
             | Self::ABCMeta
             | Self::Super
             | Self::StdlibAlias
@@ -669,6 +687,9 @@ impl KnownClass {
             | KnownClass::Member
             | KnownClass::Nonmember
             | KnownClass::StrEnum
+            | KnownClass::IntEnum
+            | KnownClass::Flag
+            | KnownClass::IntFlag
             | KnownClass::ABCMeta
             | KnownClass::GenericAlias
             | KnownClass::ModuleType
@@ -796,6 +817,9 @@ impl KnownClass {
             Self::Member => "member",
             Self::Nonmember => "nonmember",
             Self::StrEnum => "StrEnum",
+            Self::IntEnum => "IntEnum",
+            Self::Flag => "Flag",
+            Self::IntFlag => "IntFlag",
             Self::ABCMeta => "ABCMeta",
             Self::Super => "super",
             Self::Iterable => "Iterable",
@@ -1126,7 +1150,10 @@ impl KnownClass {
             | Self::Auto
             | Self::Member
             | Self::Nonmember
-            | Self::StrEnum => KnownModule::Enum,
+            | Self::StrEnum
+            | Self::IntEnum
+            | Self::Flag
+            | Self::IntFlag => KnownModule::Enum,
             Self::GenericAlias
             | Self::ModuleType
             | Self::FunctionType
@@ -1282,6 +1309,9 @@ impl KnownClass {
             | Self::Member
             | Self::Nonmember
             | Self::StrEnum
+            | Self::IntEnum
+            | Self::Flag
+            | Self::IntFlag
             | Self::ABCMeta
             | Self::Super
             | Self::NewType
@@ -1377,6 +1407,9 @@ impl KnownClass {
             | Self::Member
             | Self::Nonmember
             | Self::StrEnum
+            | Self::IntEnum
+            | Self::Flag
+            | Self::IntFlag
             | Self::ABCMeta
             | Self::Super
             | Self::UnionType
@@ -1478,6 +1511,9 @@ impl KnownClass {
             "StrEnum" if Program::get(db).python_version(db) >= PythonVersion::PY311 => {
                 &[Self::StrEnum]
             }
+            "IntEnum" => &[Self::IntEnum],
+            "Flag" => &[Self::Flag],
+            "IntFlag" => &[Self::IntFlag],
             "auto" => &[Self::Auto],
             "member" => &[Self::Member],
             "nonmember" => &[Self::Nonmember],
@@ -1564,6 +1600,9 @@ impl KnownClass {
             | Self::Member
             | Self::Nonmember
             | Self::StrEnum
+            | Self::IntEnum
+            | Self::Flag
+            | Self::IntFlag
             | Self::ABCMeta
             | Self::Super
             | Self::NotImplementedType

crates/ty_python_semantic/src/types/enums.rs

@@ -424,8 +424,9 @@ pub(crate) fn enum_metadata<'db>(
     })
 }
 
-/// Iterates over parent enum classes in the MRO, skipping known classes
-/// (like `Enum`, `StrEnum`, etc.) that we handle specially.
+/// Iterates over parent enum classes in the MRO, skipping known enum
+/// infrastructure classes but including `IntEnum`, `Flag`, and `IntFlag`
+/// which declare `_value_` annotations that should be inherited.
 fn iter_parent_enum_classes<'db>(
     db: &'db dyn Db,
     class: StaticClassLiteral<'db>,
@@ -436,7 +437,13 @@ fn iter_parent_enum_classes<'db>(
         .filter_map(ClassBase::into_class)
         .filter_map(move |class_type| {
             let base = class_type.class_literal(db).as_static()?;
-            (base.known(db).is_none() && is_enum_class_by_inheritance(db, base)).then_some(base)
+            let is_traversable = base.known(db).is_none_or(|k| {
+                matches!(
+                    k,
+                    KnownClass::IntEnum | KnownClass::Flag | KnownClass::IntFlag
+                )
+            });
+            (is_traversable && is_enum_class_by_inheritance(db, base)).then_some(base)
         })
 }
 

crates/ty_python_semantic/src/types/infer/builder/enum_call.rs

@@ -17,8 +17,16 @@ pub(crate) fn enum_functional_call_base<'db>(db: &'db dyn Db, ty: Type<'db>) ->
     let ClassLiteral::Static(cls) = ty.as_class_literal()? else {
         return None;
     };
-    cls.known(db)
-        .filter(|k| matches!(k, KnownClass::Enum | KnownClass::StrEnum))
+    cls.known(db).filter(|k| {
+        matches!(
+            k,
+            KnownClass::Enum
+                | KnownClass::StrEnum
+                | KnownClass::IntEnum
+                | KnownClass::Flag
+                | KnownClass::IntFlag
+        )
+    })
 }
 
 impl<'db> TypeInferenceBuilder<'db, '_> {
@@ -32,6 +40,19 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
         let args = &call_expr.arguments.args;
         let keywords = &call_expr.arguments.keywords;
 
+        // bail out on unknown keywords so normal overload resolution can diagnose them
+        let has_unknown_keyword = keywords.iter().any(|kw| {
+            kw.arg.as_ref().is_some_and(|name| {
+                !matches!(
+                    name.as_str(),
+                    "value" | "names" | "start" | "type" | "module" | "qualname" | "boundary"
+                )
+            })
+        });
+        if has_unknown_keyword {
+            return None;
+        }
+
         let value_kw = call_expr.arguments.find_keyword("value");
         let names_kw = call_expr.arguments.find_keyword("names");
         let start_kw = call_expr.arguments.find_keyword("start");
@@ -172,13 +193,33 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
         for (i, elt) in elts.iter().enumerate() {
             let ty = self.expression_type(elt);
             if let Some(string_lit) = ty.as_string_literal() {
+                // string element: auto-assign incrementing int value
                 let member_name = Name::new(string_lit.value(db));
                 let v = Type::int_literal(start + i64::try_from(i).unwrap_or(0));
                 members.push((member_name, v));
+            } else if let Some((name, value)) = self.extract_enum_tuple_entry(elt) {
+                members.push((name, value));
             } else {
                 return (vec![], false);
             }
         }
         (members, true)
     }
+
+    /// Extract a `(name, value)` pair from a tuple element like `("RED", 1)`.
+    fn extract_enum_tuple_entry(&mut self, elt: &ast::Expr) -> Option<(Name, Type<'db>)> {
+        let pair = match elt {
+            ast::Expr::Tuple(tup) => &tup.elts,
+            ast::Expr::List(list) => &list.elts,
+            _ => return None,
+        };
+        let [name_expr, value_expr] = &**pair else {
+            return None;
+        };
+        let db = self.db();
+        let name_ty = self.expression_type(name_expr);
+        let name = Name::new(name_ty.as_string_literal()?.value(db));
+        let value = self.expression_type(value_expr);
+        Some((name, value))
+    }
 }