fix Support __class_getitem__ (#1349)

Summary:
fix #256

Pull Request resolved: #1349

Reviewed By: rchen152

Differential Revision: D85180587

Pulled By: stroxler

fbshipit-source-id: ef412c8025a5b85eed0bec016d147b52780162cb

Author: asukaminato0721

crates/pyrefly_python/src/dunder.rs

@@ -13,6 +13,7 @@ pub const AEXIT: Name = Name::new_static("__aexit__");
 pub const ALL: Name = Name::new_static("__all__");
 pub const BOOL: Name = Name::new_static("__bool__");
 pub const CALL: Name = Name::new_static("__call__");
+pub const CLASS_GETITEM: Name = Name::new_static("__class_getitem__");
 pub const CONTAINS: Name = Name::new_static("__contains__");
 pub const DATACLASS_FIELDS: Name = Name::new_static("__dataclass_fields__");
 pub const DELATTR: Name = Name::new_static("__delattr__");

pyrefly/lib/alt/expr.rs

@@ -1855,12 +1855,44 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                         )
                     }
                 }
-                Type::ClassDef(cls) => Type::type_form(self.specialize(
-                    &cls,
-                    xs.map(|x| self.expr_untype(x, TypeFormContext::TypeArgument, errors)),
-                    range,
-                    errors,
-                )),
+                Type::ClassDef(cls) => {
+                    let metadata = self.get_metadata_for_class(&cls);
+                    let class_getitem_result = if self.get_class_tparams(&cls).is_empty()
+                        && !metadata.has_base_any()
+                        && !metadata.is_new_type()
+                    {
+                        let class_ty = Type::ClassDef(cls.dupe());
+                        // TODO(stroxler): Add a new API, similar to `type_of_attr_get` but returning a
+                        // LookupResult or an Optional type, that we could use here to avoid the double lookup.
+                        if self.has_attr(&class_ty, &dunder::CLASS_GETITEM) {
+                            let cls_value = self.promote_silently(&cls);
+                            let call_args = [CallArg::ty(&cls_value, range), CallArg::expr(slice)];
+                            Some(self.call_method_or_error(
+                                &class_ty,
+                                &dunder::CLASS_GETITEM,
+                                range,
+                                &call_args,
+                                &[],
+                                errors,
+                                Some(&|| ErrorContext::Index(self.for_display(class_ty.clone()))),
+                            ))
+                        } else {
+                            None
+                        }
+                    } else {
+                        None
+                    };
+                    if let Some(result) = class_getitem_result {
+                        result
+                    } else {
+                        Type::type_form(self.specialize(
+                            &cls,
+                            xs.map(|x| self.expr_untype(x, TypeFormContext::TypeArgument, errors)),
+                            range,
+                            errors,
+                        ))
+                    }
+                }
                 Type::Type(box Type::SpecialForm(special)) => {
                     self.apply_special_form(special, slice, range, errors)
                 }

pyrefly/lib/test/simple.rs

@@ -1737,6 +1737,19 @@ def f(condition: bool):
     "#,
 );
 
+testcase!(
+    test_class_getitem_magic_dunder,
+    r#"
+from typing import assert_type
+
+class Foo:
+    def __class_getitem__(cls, item: int) -> str:
+        return str(item)
+
+assert_type(Foo[0], str)
+"#,
+);
+
 testcase!(test_panic_docstring, "\"\"\" F\n\u{85}\"\"\"",);
 
 testcase!(