Support iterating over a class object

Reviewed By: yangdanny97

Differential Revision: D67615541

fbshipit-source-id: f227cefb10b03883c48bc6a1c1e5fdc270dcdd5e

Author: rchen152

pyre2/conformance/third_party/conformance.exp

@@ -8326,26 +8326,6 @@
     }
   ],
   "enums_behaviors.py": [
-    {
-      "code": -2,
-      "column": 14,
-      "concise_description": "TODO: Answers::solve_binding - Binding::IterableValue",
-      "description": "TODO: Answers::solve_binding - Binding::IterableValue",
-      "line": 19,
-      "name": "PyreError",
-      "stop_column": 19,
-      "stop_line": 19
-    },
-    {
-      "code": -2,
-      "column": 5,
-      "concise_description": "assert_type(Any, Color) failed",
-      "description": "assert_type(Any, Color) failed",
-      "line": 20,
-      "name": "PyreError",
-      "stop_column": 30,
-      "stop_line": 20
-    },
     {
       "code": -2,
       "column": 1,

pyre2/conformance/third_party/conformance.result

@@ -603,8 +603,6 @@
   "directives_version_platform.py": [],
   "enums_behaviors.py": [
     "Line 39: Expected 1 errors",
-    "Line 19: Unexpected errors ['TODO: Answers::solve_binding - Binding::IterableValue']",
-    "Line 20: Unexpected errors ['assert_type(Any, Color) failed']",
     "Line 26: Unexpected errors ['assert_type(type[Color], Color) failed', 'Expected 0 type arguments for class `Color`, got 1.', \"untype, got Literal['RED']\"]",
     "Line 27: Unexpected errors ['assert_type(Enum, Color) failed']"
   ],

pyre2/conformance/third_party/results.json

@@ -3,7 +3,7 @@
   "pass": 8,
   "fail": 125,
   "pass_rate": 0.06,
-  "differences": 1440,
+  "differences": 1438,
   "passing": [
     "annotations_coroutines.py",
     "directives_no_type_check.py",
@@ -59,7 +59,7 @@
     "directives_reveal_type.py": 4,
     "directives_type_checking.py": 3,
     "directives_type_ignore_file1.py": 1,
-    "enums_behaviors.py": 5,
+    "enums_behaviors.py": 3,
     "enums_definition.py": 1,
     "enums_expansion.py": 5,
     "enums_member_values.py": 2,

pyre2/pyre2/bin/alt/answers.rs

@@ -66,6 +66,7 @@ use crate::types::annotation::Qualifier;
 use crate::types::callable::Param;
 use crate::types::callable::Required;
 use crate::types::class::Class;
+use crate::types::class::ClassType;
 use crate::types::class_metadata::ClassMetadata;
 use crate::types::literal::Lit;
 use crate::types::module::Module;
@@ -671,6 +672,19 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         }
     }
 
+    fn iterate_by_metaclass(&self, cls: &ClassType, range: TextRange) -> Iterable {
+        let call_method =
+            |method_name: &Name, range: TextRange, args: &[CallArg], keywords: &[Keyword]| {
+                self.call_metaclass_method(cls, method_name, range, args, keywords)
+            };
+        let ty = self
+            .iterate_by_method(&call_method, range)
+            .unwrap_or_else(|| {
+                self.error(range, format!("Class object `{}` is not iterable", cls))
+            });
+        Iterable::OfType(ty)
+    }
+
     fn iterate(&self, iterable: &Type, range: TextRange) -> Iterable {
         match iterable {
             Type::ClassType(cls) => {
@@ -695,6 +709,10 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             Type::Var(v) if let Some(_guard) = self.recurser.recurse(*v) => {
                 self.iterate(&self.solver().force_var(*v), range)
             }
+            Type::ClassDef(cls) => {
+                self.iterate_by_metaclass(&self.promote_to_class_type(cls, range), range)
+            }
+            Type::Type(box Type::ClassType(cls)) => self.iterate_by_metaclass(cls, range),
             _ => Iterable::OfType(
                 self.error_todo("Answers::solve_binding - Binding::IterableValue", range),
             ),

pyre2/pyre2/bin/alt/expr.rs

@@ -267,7 +267,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
 
     /// Calls a method on the metaclass of `cls`. Returns None without attempting a call if we
     /// don't find a method defined on a metaclass that isn't the default `builtins.type`.
-    fn call_metaclass_method(
+    pub fn call_metaclass_method(
         &self,
         cls: &ClassType,
         method_name: &Name,

pyre2/pyre2/bin/test/enums.rs

@@ -20,3 +20,16 @@ class MyEnum(Enum):
 assert_type(MyEnum.X, Literal[MyEnum.X])
 "#,
 );
+
+testcase!(
+    test_iterate,
+    r#"
+from typing import assert_type
+from enum import Enum
+class E(Enum):
+    X = 1
+    Y = 2
+for e in E:
+    assert_type(e, E)
+    "#,
+);

pyre2/pyre2/bin/test/simple.rs

@@ -747,6 +747,33 @@ def f(x: A):
     "#,
 );
 
+testcase!(
+    test_iterable_class_error,
+    r#"
+class A:
+    pass
+for _ in A:  # E: Class object `A` is not iterable
+    pass
+    "#,
+);
+
+testcase!(
+    test_iterable_generic_class,
+    r#"
+from typing import Iterator, assert_type
+class M(type):
+    def __iter__(self) -> Iterator[int]: ...
+class A[T](metaclass=M):
+    pass
+class B[T]:
+    pass
+for x in A[str]:
+    assert_type(x, int)
+for _ in B[str]:  # E: Class object `B[str]` is not iterable
+    pass
+    "#,
+);
+
 testcase!(
     test_iterable_bad_callable,
     r#"

pyre2/pyre2/bin/test/stdlib.rs

@@ -11,7 +11,9 @@
 use crate::module::module_name::ModuleName;
 
 static ENUM: &str = r#"
-class Enum: ...
+class EnumMeta(type):
+    def __iter__[EnumMemberT](self: type[EnumMemberT]) -> Iterator[EnumMemberT]: ...
+class Enum(metaclass=EnumMeta): ...
 class StrEnum(str, Enum): ...
 class IntEnum(int, Enum): ...
 "#;