From a7c5cdc4a5554838ae360025a0f76b7131f12438 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 0d511d5d1a..a12fefeb09 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -52,6 +52,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; @@ -75,6 +76,7 @@ use crate::binding::binding::KeyDecoratedFunction; use crate::binding::binding::KeyDecorator; use crate::binding::binding::KeyLegacyTypeParam; use crate::binding::classmethod::MethodReceiverKind; +use crate::binding::pytest::is_pytest_parametrize_decorator; use crate::config::error_kind::ErrorKind; use crate::error::collector::ErrorCollector; use crate::error::context::TypeCheckContext; @@ -215,6 +217,7 @@ impl ParentParamHints { #[derive(Clone, Debug)] struct DecoratorParamHints { positional: Vec, + by_name: SmallMap, next_positional: usize, } @@ -245,6 +248,7 @@ impl DecoratorParamHints { } else { Some(Self { positional, + by_name: SmallMap::new(), next_positional: 0, }) } @@ -253,6 +257,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 @@ -282,6 +319,128 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, 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, idx: Idx, @@ -524,6 +683,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { calls_super_method, ..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); @@ -575,6 +735,13 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, 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( @@ -1439,7 +1606,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, 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 { @@ -1473,7 +1640,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, 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 { @@ -1559,6 +1726,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, 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)); @@ -1571,7 +1741,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { x.default.as_deref(), is_stub, 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 8dda2c21ce..ed6f8e9c95 100644 --- a/pyrefly/lib/binding/pytest.rs +++ b/pyrefly/lib/binding/pytest.rs @@ -195,6 +195,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()) + ) +} + pub(crate) 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 7e5033a813..ec2d9a718c 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -2074,6 +2074,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 ab3956120fd86a61cb2656c3ff056c7ea38b6c8a 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 ec2d9a718c..fbda1c3ae6 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -2095,7 +2095,7 @@ testcase!( env_pytest_parametrize(), r#" import pytest -from typing import reveal_type +from typing import assert_type @pytest.mark.parametrize( ("exception", "expected_error"), @@ -2106,10 +2106,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) "#, );