From 1ad195c0431790d51a4e596b1972c08f2d7494d1 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 30 Jun 2026 02:39:54 +0900 Subject: [PATCH 1/2] can infer tuple[T1 | T3, T2 | T4] --- pyrefly/lib/alt/function.rs | 176 ++++++++++++++++++++++++++++++++- pyrefly/lib/binding/pytest.rs | 20 ++++ pyrefly/lib/test/decorators.rs | 39 ++++++++ 3 files changed, 232 insertions(+), 3 deletions(-) diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index 6a266da643..4271ef2a87 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -44,6 +44,7 @@ use ruff_python_ast::UnaryOp; use ruff_python_ast::name::Name; use ruff_text_size::Ranged; use ruff_text_size::TextRange; +use starlark_map::small_map::Entry; use starlark_map::small_map::SmallMap; use starlark_map::small_set::SmallSet; use vec1::Vec1; @@ -67,6 +68,7 @@ use crate::binding::binding::KeyClass; use crate::binding::binding::KeyClassMetadata; use crate::binding::binding::KeyDecorator; use crate::binding::binding::KeyLegacyTypeParam; +use crate::binding::pytest::is_pytest_parametrize_decorator; use crate::config::error_kind::ErrorKind; use crate::error::collector::ErrorCollector; use crate::error::context::TypeCheckContext; @@ -206,6 +208,7 @@ impl ParentParamHints { #[derive(Clone, Debug)] struct DecoratorParamHints { positional: Vec, + by_name: SmallMap, next_positional: usize, } @@ -236,6 +239,7 @@ impl DecoratorParamHints { } else { Some(Self { positional, + by_name: SmallMap::new(), next_positional: 0, }) } @@ -244,6 +248,39 @@ impl DecoratorParamHints { } } + fn from_named(by_name: SmallMap) -> Option { + if by_name.is_empty() { + None + } else { + Some(Self { + positional: Vec::new(), + by_name, + next_positional: 0, + }) + } + } + + fn merge(&mut self, other: Self, solver: &AnswersSolver<'_, Ans>) { + self.positional.extend(other.positional); + for (name, ty) in other.by_name { + match self.by_name.entry(name) { + Entry::Occupied(mut entry) => { + let old = entry.get().clone(); + *entry.get_mut() = solver.union(old, ty); + } + Entry::Vacant(entry) => { + entry.insert(ty); + } + } + } + } + + fn take(&mut self, name: &Identifier) -> Option { + self.by_name + .shift_remove(&name.id) + .or_else(|| self.next_positional()) + } + fn next_positional(&mut self) -> Option { if self.next_positional >= self.positional.len() { None @@ -268,6 +305,128 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { }) } + fn pytest_parametrize_param_hints( + &self, + decorators: &[Idx], + errors: &ErrorCollector, + ) -> Option { + let aliases = self.bindings().pytest_info()?.aliases(); + let mut by_name: SmallMap = SmallMap::new(); + for key in decorators { + let binding = self.bindings().get::(*key); + if !is_pytest_parametrize_decorator(&binding.expr, aliases) { + continue; + } + let Expr::Call(call) = &binding.expr else { + unreachable!("pytest parametrize decorator is a call") + }; + let [argnames, argvalues, ..] = &*call.arguments.args else { + continue; + }; + let Some(names) = Self::pytest_parametrize_names(argnames) else { + continue; + }; + let Some(types) = self.pytest_parametrize_value_types(argvalues, names.len(), errors) + else { + continue; + }; + for (name, ty) in names.into_iter().zip(types) { + match by_name.entry(name) { + Entry::Occupied(mut entry) => { + let old = entry.get().clone(); + *entry.get_mut() = self.union(old, ty); + } + Entry::Vacant(entry) => { + entry.insert(ty); + } + } + } + } + DecoratorParamHints::from_named(by_name) + } + + fn pytest_parametrize_names(expr: &Expr) -> Option> { + match expr { + Expr::StringLiteral(lit) => { + let names = lit + .value + .to_str() + .split(',') + .map(str::trim) + .filter(|name| !name.is_empty()) + .map(Name::new) + .collect::>(); + (!names.is_empty()).then_some(names) + } + Expr::Tuple(tuple) => Self::pytest_parametrize_names_from_elts(&tuple.elts), + Expr::List(list) => Self::pytest_parametrize_names_from_elts(&list.elts), + _ => None, + } + } + + fn pytest_parametrize_names_from_elts(elts: &[Expr]) -> Option> { + let mut names = Vec::with_capacity(elts.len()); + for elt in elts { + let Expr::StringLiteral(lit) = elt else { + return None; + }; + let name = lit.value.to_str(); + if name.is_empty() { + return None; + } + names.push(Name::new(name)); + } + Some(names) + } + + fn pytest_parametrize_value_types( + &self, + expr: &Expr, + arity: usize, + errors: &ErrorCollector, + ) -> Option> { + let rows = match expr { + Expr::List(list) => list.elts.as_slice(), + Expr::Tuple(tuple) => tuple.elts.as_slice(), + _ => return None, + }; + let hint_errors = ErrorCollector::new(errors.module().dupe(), errors.style()); + let mut columns = vec![Vec::new(); arity]; + for row in rows { + if arity == 1 { + columns[0].push( + self.expr_infer(row, &hint_errors) + .promote_implicit_literals(self.stdlib), + ); + continue; + } + let elts = match row { + Expr::Tuple(tuple) => tuple.elts.as_slice(), + Expr::List(list) => list.elts.as_slice(), + _ => return None, + }; + if elts.len() != arity { + return None; + } + for (column, elt) in columns.iter_mut().zip(elts) { + column.push( + self.expr_infer(elt, &hint_errors) + .promote_implicit_literals(self.stdlib), + ); + } + } + columns + .into_iter() + .map(|column| { + if column.is_empty() { + None + } else { + Some(self.unions(column)) + } + }) + .collect() + } + pub fn solve_function_binding( &self, def: DecoratedFunction, @@ -467,6 +626,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { is_return_inferred, ..Default::default() }; + let pytest_param_hints = self.pytest_parametrize_param_hints(decorators, errors); let mut found_class_property = false; let decorators = Box::from_iter(decorators.iter().filter_map(|k| { let decorator = self.get_idx(*k); @@ -504,6 +664,13 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { })); let mut decorator_param_hints = self.decorator_param_hints(&decorators); + if let Some(pytest_hints) = pytest_param_hints { + if let Some(hints) = &mut decorator_param_hints { + hints.merge(pytest_hints, self); + } else { + decorator_param_hints = Some(pytest_hints); + } + } let mut parent_param_hints = if flags.is_override { defining_cls.as_ref().and_then(|cls| { self.inherited_method_signature(cls, &def.name.id).and_then( @@ -1172,7 +1339,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { params.extend(def.parameters.posonlyargs.iter().map(|x| { let decorator_hint = decorator_param_hints .as_mut() - .and_then(|hint| hint.next_positional()); + .and_then(|hint| hint.take(&x.parameter.name)); let parent_hint = if self_type.is_some() { None } else { @@ -1206,7 +1373,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { params.extend(def.parameters.args.iter().map(|x| { let decorator_hint = decorator_param_hints .as_mut() - .and_then(|hint| hint.next_positional()); + .and_then(|hint| hint.take(&x.parameter.name)); let parent_hint = if self_type.is_some() { None } else { @@ -1292,6 +1459,9 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { ); } params.extend(def.parameters.kwonlyargs.iter().map(|x| { + let decorator_hint = decorator_param_hints + .as_mut() + .and_then(|hint| hint.by_name.shift_remove(&x.parameter.name.id)); let parent_hint = parent_param_hints .as_mut() .and_then(|hint| hint.take_kwonly(&x.parameter.name)); @@ -1304,7 +1474,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { x.default.as_deref(), stub_or_impl, self_type, - parent_hint, + decorator_hint.or(parent_hint), errors, ); if is_unannotated { diff --git a/pyrefly/lib/binding/pytest.rs b/pyrefly/lib/binding/pytest.rs index f324eab2a9..3dae21f857 100644 --- a/pyrefly/lib/binding/pytest.rs +++ b/pyrefly/lib/binding/pytest.rs @@ -169,6 +169,26 @@ pub fn is_pytest_fixture_function(function_def: &StmtFunctionDef, aliases: &Pyte .any(|decorator| is_pytest_fixture_decorator(&decorator.expression, aliases)) } +pub fn is_pytest_parametrize_decorator(expr: &Expr, aliases: &PytestAliases) -> bool { + let Expr::Call(call) = expr else { + return false; + }; + let Expr::Attribute(parametrize) = call.func.as_ref() else { + return false; + }; + if parametrize.attr.id() != "parametrize" { + return false; + } + let Expr::Attribute(mark) = parametrize.value.as_ref() else { + return false; + }; + mark.attr.id() == "mark" + && matches!( + mark.value.as_ref(), + Expr::Name(base) if aliases.is_pytest_module_alias(base.id()) + ) +} + fn is_pytest_fixture_decorator(expr: &Expr, aliases: &PytestAliases) -> bool { match expr { Expr::Name(name) => aliases.is_fixture_alias(name.id()), diff --git a/pyrefly/lib/test/decorators.rs b/pyrefly/lib/test/decorators.rs index 1a3ce44c71..4006f53a19 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -730,6 +730,45 @@ assert_type(r, int) "#, ); +fn env_pytest_parametrize() -> TestEnv { + TestEnv::one_with_path( + "pytest", + "pytest.pyi", + r#" +class MarkDecorator: + def __call__[T](self, func: T) -> T: ... + +class MarkGenerator: + def parametrize(self, argnames: object, argvalues: object) -> MarkDecorator: ... + +mark: MarkGenerator +"#, + ) +} + +testcase!( + test_pytest_parametrize_infers_parameter_types, + env_pytest_parametrize(), + r#" +import pytest +from typing import reveal_type + +@pytest.mark.parametrize( + ("exception", "expected_error"), + [ + (Exception("OAuth error"), "OAuth process failed"), + (ValueError("Invalid token"), "OAuth process failed"), + (KeyError("Missing key"), "OAuth process failed"), + ], +) +def test_oauth_error(exception, expected_error) -> None: + reveal_type(exception) # E: revealed type: Exception | KeyError | ValueError + reveal_type(expected_error) # E: revealed type: str + +reveal_type(test_oauth_error) # E: revealed type: (exception: Exception | KeyError | ValueError, expected_error: str) -> None +"#, +); + fn env_numba() -> TestEnv { let mut env = TestEnv::one_with_path( "numba", From 9b2d2adde19f39a255bbe5efe562e8ef05d57808 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 30 Jun 2026 02:56:14 +0900 Subject: [PATCH 2/2] assert --- pyrefly/lib/test/decorators.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/pyrefly/lib/test/decorators.rs b/pyrefly/lib/test/decorators.rs index 4006f53a19..cdfc7a8272 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -751,7 +751,7 @@ testcase!( env_pytest_parametrize(), r#" import pytest -from typing import reveal_type +from typing import assert_type @pytest.mark.parametrize( ("exception", "expected_error"), @@ -762,10 +762,8 @@ from typing import reveal_type ], ) def test_oauth_error(exception, expected_error) -> None: - reveal_type(exception) # E: revealed type: Exception | KeyError | ValueError - reveal_type(expected_error) # E: revealed type: str - -reveal_type(test_oauth_error) # E: revealed type: (exception: Exception | KeyError | ValueError, expected_error: str) -> None + assert_type(exception, Exception | KeyError | ValueError) + assert_type(expected_error, str) "#, );