diff --git a/pyrefly/lib/alt/class/class_metadata.rs b/pyrefly/lib/alt/class/class_metadata.rs index 4ecf35effe..896eb1c209 100644 --- a/pyrefly/lib/alt/class/class_metadata.rs +++ b/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], 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() diff --git a/pyrefly/lib/binding/binding.rs b/pyrefly/lib/binding/binding.rs index 4a6424e05e..8ac87d9d7a 100644 --- a/pyrefly/lib/binding/binding.rs +++ b/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]>, /// Is this a new type? True only for synthesized classes created from a `NewType` call. diff --git a/pyrefly/lib/binding/class.rs b/pyrefly/lib/binding/class.rs index a0098e7a80..841f5d63f5 100644 --- a/pyrefly/lib/binding/class.rs +++ b/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, - keywords: Box<[(Name, Expr)]>, + keywords: Box<[Keyword]>, // name, position, annotation, value member_definitions: Vec<(String, TextRange, Option, Option)>, 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}`") diff --git a/pyrefly/lib/test/class_keywords.rs b/pyrefly/lib/test/class_keywords.rs index d34757e694..8d8cdd12ca 100644 --- a/pyrefly/lib/test/class_keywords.rs +++ b/pyrefly/lib/test/class_keywords.rs @@ -52,6 +52,7 @@ 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!( @@ -59,6 +60,8 @@ class A(foo=True): pass 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 "#, );