fix Unpacking dict in class header (i.e. in class kwargs) #1762

pyrefly/lib/alt/class/class_metadata.rs

@@ -9,7 +9,6 @@ use std::iter;
 use std::sync::Arc;
 
 use dupe::Dupe;
-use itertools::Either;
 use itertools::Itertools;
 use pyrefly_graph::index::Idx;
 use pyrefly_python::dunder;
@@ -28,6 +27,7 @@ use pyrefly_util::display::DisplayWithCtx;
 use pyrefly_util::prelude::SliceExt;
 use pyrefly_util::prelude::VecExt;
 use ruff_python_ast::Expr;
+use ruff_python_ast::Keyword;
 use ruff_python_ast::name::Name;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
@@ -138,11 +138,47 @@ pub(crate) struct TransformDataclass {
 }
 
 impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
+    fn extend_unpacked_class_keywords(
+        &self,
+        keyword: &Keyword,
+        keywords: &mut Vec<(Name, Annotation)>,
+        errors: &ErrorCollector,
+    ) {
+        let ty = self.expr_infer(&keyword.value, errors);
+        if let Type::TypedDict(typed_dict) = ty {
+            for (name, field) in self.typed_dict_fields(&typed_dict) {
+                keywords.push((name, Annotation::new_type(field.ty)));
+            }
+        } else if let Some((key, _)) = self.unwrap_mapping(&ty) {
+            if !self.is_subset_eq(&key, &self.heap.mk_class_type(self.stdlib.str().clone())) {
+                self.error(
+                    errors,
+                    keyword.value.range(),
+                    ErrorKind::BadUnpacking,
+                    format!(
+                        "Expected argument after ** to have `str` keys, got: {}",
+                        self.for_display(key)
+                    ),
+                );
+            }
+        } else {
+            self.error(
+                errors,
+                keyword.value.range(),
+                ErrorKind::BadUnpacking,
+                format!(
+                    "Expected argument after ** to be a mapping, got: {}",
+                    self.for_display(ty)
+                ),
+            );
+        }
+    }
+
     pub fn class_metadata_of(
         &self,
         cls: &Class,
         bases: &[BaseClass],
-        keywords: &[(Name, Expr)],
+        raw_keywords: &[Keyword],
         decorators: &[Idx<KeyDecorator>],
         is_new_type: bool,
         pydantic_config_dict: &PydanticConfigDict,
@@ -183,11 +219,18 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         let bases_with_metadata = self.bases_with_metadata(parsed_results, is_new_type, errors);
 
         // Compute class keywords, including the metaclass.
-        let (metaclasses, keywords): (Vec<_>, Vec<(_, _)>) =
-            keywords.iter().partition_map(|(n, x)| match n.as_str() {
-                "metaclass" => Either::Left(x),
-                _ => Either::Right((n.clone(), self.expr_class_keyword(x, errors))),
-            });
+        let mut metaclasses = Vec::new();
+        let mut keywords = Vec::new();
+        for keyword in raw_keywords {
+            match &keyword.arg {
+                Some(name) if name.id == "metaclass" => metaclasses.push(&keyword.value),
+                Some(name) => keywords.push((
+                    name.id.clone(),
+                    self.expr_class_keyword(&keyword.value, errors),
+                )),
+                None => self.extend_unpacked_class_keywords(keyword, &mut keywords, errors),
+            }
+        }
 
         let base_metaclasses = bases_with_metadata
             .iter()

pyrefly/lib/binding/binding.rs

@@ -40,6 +40,7 @@ use ruff_python_ast::ExprSubscript;
 use ruff_python_ast::ExprYield;
 use ruff_python_ast::ExprYieldFrom;
 use ruff_python_ast::Identifier;
+use ruff_python_ast::Keyword;
 use ruff_python_ast::Parameters;
 use ruff_python_ast::StmtAugAssign;
 use ruff_python_ast::StmtClassDef;
@@ -3207,7 +3208,7 @@ pub struct BindingClassMetadata {
     /// The class keywords (these are keyword args that appear in the base class list, the
     /// Python runtime will dispatch most of them to the metaclass, but the metaclass
     /// itself can also potentially be one of these).
-    pub keywords: Box<[(Name, Expr)]>,
+    pub keywords: Box<[Keyword]>,
     /// The class decorators.
     pub decorators: Box<[Idx<KeyDecorator>]>,
     /// Is this a new type? True only for synthesized classes created from a `NewType` call.

pyrefly/lib/binding/class.rs

@@ -386,16 +386,8 @@ impl<'a> BindingsBuilder<'a> {
         let mut keywords = Vec::new();
         if let Some(args) = &mut x.arguments {
             args.keywords.iter_mut().for_each(|keyword| {
-                if let Some(name) = &keyword.arg {
-                    self.ensure_expr(&mut keyword.value, class_object.usage());
-                    keywords.push((name.id.clone(), keyword.value.clone()));
-                } else {
-                    self.error(
-                        keyword.range(),
-                        ErrorKind::InvalidInheritance,
-                        "Unpacking is not supported in class header".to_owned(),
-                    )
-                }
+                self.ensure_expr(&mut keyword.value, class_object.usage());
+                keywords.push(keyword.clone());
             });
         }
         let bases: Arc<[BaseClass]> = Arc::from(bases.into_boxed_slice());
@@ -1143,7 +1135,7 @@ impl<'a> BindingsBuilder<'a> {
         class_indices: ClassIndices,
         parent: &NestingContext,
         base: Option<Expr>,
-        keywords: Box<[(Name, Expr)]>,
+        keywords: Box<[Keyword]>,
         // name, position, annotation, value
         member_definitions: Vec<(String, TextRange, Option<Expr>, Option<ExprOrBinding>)>,
         illegal_identifier_handling: IllegalIdentifierHandling,
@@ -1591,8 +1583,8 @@ impl<'a> BindingsBuilder<'a> {
                 (Some(name), _) if name == "extra_items" => Some(name),
                 _ => None,
             };
-            if let Some(kw_name) = recognized_kw {
-                base_class_keywords.push((kw_name.clone(), kw.value.clone()));
+            if recognized_kw.is_some() {
+                base_class_keywords.push(kw.clone());
             } else {
                 let msg = if let Some(name) = &kw.arg {
                     format!("Unrecognized keyword argument `{name}`")

pyrefly/lib/test/class_keywords.rs

@@ -52,13 +52,16 @@ fn test_look_up_class_keywords() {
     let (handle, state) = mk_state(
         r#"
 class A(foo=True): pass
+class B(**{"foo": True, "bar": 1}): pass
 "#,
     );
     assert_eq!(
         get_class_keywords("A", "foo", &handle, &state),
         vec![Lit::Bool(true).to_implicit_type()],
     );
     assert_eq!(get_class_keywords("A", "bar", &handle, &state), vec![]);
+    assert_eq!(get_class_keywords("B", "foo", &handle, &state).len(), 1);
+    assert_eq!(get_class_keywords("B", "bar", &handle, &state).len(), 1);
 }
 
 #[test]
@@ -155,10 +158,19 @@ f(C2[int])
 );
 
 testcase!(
-    test_illegal_unpacking,
+    test_class_keyword_unpacking,
     r#"
-def f() -> dict: ...
-class A(**f):  # E: Unpacking is not supported in class header
+from typing import Any
+
+meta: dict[str, Any] = {}
+class A(**meta, tag="A"):
+    pass
+
+class B(**1):  # E: Expected argument after ** to be a mapping, got: Literal[1]
+    pass
+
+bad_keys: dict[int, str] = {}
+class C(**bad_keys):  # E: Expected argument after ** to have `str` keys, got: int
     pass
     "#,
 );