[ty] Support iterating over enums

crates/ty_python_semantic/resources/mdtest/enums.md

@@ -406,8 +406,7 @@ class Color(Enum):
     BLUE = 3
 
 for color in Color:
-    # TODO: Should be `Color`
-    reveal_type(color)  # revealed: Unknown
+    reveal_type(color)  # revealed: Color
 
 # TODO: Should be `list[Color]`
 reveal_type(list(Color))  # revealed: list[Unknown]

crates/ty_python_semantic/src/types/call/bind.rs

@@ -19,6 +19,7 @@ use crate::types::diagnostic::{
     NO_MATCHING_OVERLOAD, PARAMETER_ALREADY_ASSIGNED, TOO_MANY_POSITIONAL_ARGUMENTS,
     UNKNOWN_ARGUMENT,
 };
+use crate::types::enums::is_enum_class;
 use crate::types::function::{
     DataclassTransformerParams, FunctionDecorators, FunctionType, KnownFunction, OverloadLiteral,
 };
@@ -560,6 +561,19 @@ impl<'db> Bindings<'db> {
                         }
                     }
 
+                    // TODO: This branch can be removed once https://github.com/astral-sh/ty/issues/501 is resolved
+                    Type::BoundMethod(bound_method)
+                        if bound_method.function(db).name(db) == "__iter__"
+                            && is_enum_class(db, bound_method.self_instance(db)) =>
+                    {
+                        if let Some(enum_instance) = bound_method.self_instance(db).to_instance(db)
+                        {
+                            overload.set_return_type(
+                                KnownClass::Iterator.to_specialized_instance(db, [enum_instance]),
+                            );
+                        }
+                    }
+
                     Type::FunctionLiteral(function_type) => match function_type.known(db) {
                         Some(KnownFunction::IsEquivalentTo) => {
                             if let [Some(ty_a), Some(ty_b)] = overload.parameter_types() {

crates/ty_python_semantic/src/types/class.rs

@@ -2564,6 +2564,7 @@ pub enum KnownClass {
     NewType,
     SupportsIndex,
     Iterable,
+    Iterator,
     // Collections
     ChainMap,
     Counter,
@@ -2660,6 +2661,7 @@ impl KnownClass {
             | Self::Nonmember
             | Self::ABCMeta
             | Self::Iterable
+            | Self::Iterator
             // Empty tuples are AlwaysFalse; non-empty tuples are AlwaysTrue
             | Self::NamedTuple
             // Evaluating `NotImplementedType` in a boolean context was deprecated in Python 3.9
@@ -2753,6 +2755,7 @@ impl KnownClass {
             | Self::OrderedDict
             | Self::NewType
             | Self::Iterable
+            | Self::Iterator
             | Self::BaseExceptionGroup => false,
         }
     }
@@ -2814,6 +2817,7 @@ impl KnownClass {
             | KnownClass::NewType
             | KnownClass::SupportsIndex
             | KnownClass::Iterable
+            | KnownClass::Iterator
             | KnownClass::ChainMap
             | KnownClass::Counter
             | KnownClass::DefaultDict
@@ -2842,7 +2846,7 @@ impl KnownClass {
     /// 2. It's probably more performant.
     const fn is_protocol(self) -> bool {
         match self {
-            Self::SupportsIndex | Self::Iterable => true,
+            Self::SupportsIndex | Self::Iterable | Self::Iterator => true,
 
             Self::Any
             | Self::Bool
@@ -2968,6 +2972,7 @@ impl KnownClass {
             Self::ABCMeta => "ABCMeta",
             Self::Super => "super",
             Self::Iterable => "Iterable",
+            Self::Iterator => "Iterator",
             // For example, `typing.List` is defined as `List = _Alias()` in typeshed
             Self::StdlibAlias => "_Alias",
             // This is the name the type of `sys.version_info` has in typeshed,
@@ -3203,6 +3208,7 @@ impl KnownClass {
             | Self::NamedTuple
             | Self::StdlibAlias
             | Self::Iterable
+            | Self::Iterator
             | Self::SupportsIndex => KnownModule::Typing,
             Self::TypeAliasType
             | Self::TypeVarTuple
@@ -3311,6 +3317,7 @@ impl KnownClass {
             | Self::Field
             | Self::KwOnly
             | Self::Iterable
+            | Self::Iterator
             | Self::NamedTupleFallback => false,
         }
     }
@@ -3384,6 +3391,7 @@ impl KnownClass {
             | Self::Field
             | Self::KwOnly
             | Self::Iterable
+            | Self::Iterator
             | Self::NamedTupleFallback => false,
         }
     }
@@ -3435,6 +3443,7 @@ impl KnownClass {
             "TypeAliasType" => Self::TypeAliasType,
             "TypeVar" => Self::TypeVar,
             "Iterable" => Self::Iterable,
+            "Iterator" => Self::Iterator,
             "ParamSpec" => Self::ParamSpec,
             "ParamSpecArgs" => Self::ParamSpecArgs,
             "ParamSpecKwargs" => Self::ParamSpecKwargs,
@@ -3538,6 +3547,7 @@ impl KnownClass {
             | Self::TypeVarTuple
             | Self::NamedTuple
             | Self::Iterable
+            | Self::Iterator
             | Self::NewType => matches!(module, KnownModule::Typing | KnownModule::TypingExtensions),
             Self::Deprecated => matches!(module, KnownModule::Warnings | KnownModule::TypingExtensions),