diff --git a/pyrefly/lib/alt/jaxtyping.rs b/pyrefly/lib/alt/jaxtyping.rs index efc664fda0..980f40f711 100644 --- a/pyrefly/lib/alt/jaxtyping.rs +++ b/pyrefly/lib/alt/jaxtyping.rs @@ -191,7 +191,9 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { Binding::Import(import) => { import.module.as_str() == "jaxtyping" && import.name.as_str() == "jaxtyping" } - Binding::PossibleLegacyTParam(key, _) => self.legacy_tparam_is_jaxtyping_module(*key), + Binding::PossibleLegacyTParam(keys, _) => keys + .iter() + .any(|key| self.legacy_tparam_is_jaxtyping_module(*key)), _ => false, } } @@ -204,9 +206,9 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { | Binding::ForwardToFirstUse(idx) => { self.binding_is_jaxtyping_wrapper_origin(self.binding_following_forwards(*idx)) } - Binding::PossibleLegacyTParam(key, _) => { - self.legacy_tparam_is_jaxtyping_wrapper_origin(*key) - } + Binding::PossibleLegacyTParam(keys, _) => keys + .iter() + .any(|key| self.legacy_tparam_is_jaxtyping_wrapper_origin(*key)), _ => false, } } diff --git a/pyrefly/lib/alt/solve.rs b/pyrefly/lib/alt/solve.rs index 2ebe620dd7..f990004856 100644 --- a/pyrefly/lib/alt/solve.rs +++ b/pyrefly/lib/alt/solve.rs @@ -4909,44 +4909,64 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { #[inline(never)] fn binding_to_type_info_possible_legacy_tparam( &self, - key: Idx, + keys: &[Idx], range_if_scoped_params_exist: &Option, errors: &ErrorCollector, ) -> TypeInfo { - let ty = match &*self.get_idx(key) { - LegacyTypeParameterLookup::Parameter(p) => { - // This class or function has scoped (PEP 695) type parameters. Mixing legacy-style parameters is an error. - if let Some(r) = range_if_scoped_params_exist { - self.error( - errors, - *r, - ErrorKind::InvalidTypeVar, - format!( - "Type parameter {} is not included in the type parameter list", - self.module().display(&self.bindings().idx_to_key(key).0) - ), - ); - } - p.clone().to_value() - } - LegacyTypeParameterLookup::NotParameter(ty) => ty.clone(), - }; - match self.bindings().get(key) { - BindingLegacyTypeParam::ModuleKeyed(idx, attrs) => { - // `idx` points to a module whose attr chain may end in a legacy type - // variable that needs to be replaced with a QuantifiedValue. Since the - // binding is for the module itself, we use the mechanism for attribute - // ("facet") type narrowing to change the type produced when the final - // attr is accessed. - let module = (*self.get_idx(*idx)).clone(); - if matches!(ty, Type::QuantifiedValue(_)) { - let facets = attrs.mapped_ref(|a| FacetKind::Attribute(a.clone())); - module.with_narrow(&facets, ty) - } else { - module + let first_key = *keys + .first() + .expect("PossibleLegacyTParam bindings must always have at least one key"); + match self.bindings().get(first_key) { + BindingLegacyTypeParam::ParamKeyed(_) => { + let ty = match &*self.get_idx(first_key) { + LegacyTypeParameterLookup::Parameter(p) => { + if let Some(r) = range_if_scoped_params_exist { + self.error( + errors, + *r, + ErrorKind::InvalidTypeVar, + format!( + "Type parameter {} is not included in the type parameter list", + self.module() + .display(&self.bindings().idx_to_key(first_key).0) + ), + ); + } + p.clone().to_value() + } + LegacyTypeParameterLookup::NotParameter(ty) => ty.clone(), + }; + TypeInfo::of_ty(ty) + } + BindingLegacyTypeParam::ModuleKeyed(idx, _) => { + let mut module_info = (*self.get_idx(*idx)).clone(); + for &key in keys { + let ty = match &*self.get_idx(key) { + LegacyTypeParameterLookup::Parameter(p) => { + if let Some(r) = range_if_scoped_params_exist { + self.error( + errors, + *r, + ErrorKind::InvalidTypeVar, + format!( + "Type parameter {} is not included in the type parameter list", + self.module().display(&self.bindings().idx_to_key(key).0) + ), + ); + } + p.clone().to_value() + } + LegacyTypeParameterLookup::NotParameter(ty) => ty.clone(), + }; + if let BindingLegacyTypeParam::ModuleKeyed(_, attrs) = self.bindings().get(key) + && matches!(ty, Type::QuantifiedValue(_)) + { + let facets = attrs.mapped_ref(|a| FacetKind::Attribute(a.clone())); + module_info = module_info.with_narrow(&facets, ty); + } } + module_info } - BindingLegacyTypeParam::ParamKeyed(_) => TypeInfo::of_ty(ty), } } @@ -5015,7 +5035,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { } Binding::PossibleLegacyTParam(key, range_if_scoped_params_exist) => self .binding_to_type_info_possible_legacy_tparam( - *key, + key, range_if_scoped_params_exist, errors, ), diff --git a/pyrefly/lib/binding/binding.rs b/pyrefly/lib/binding/binding.rs index 12f5e14446..966b0b68f6 100644 --- a/pyrefly/lib/binding/binding.rs +++ b/pyrefly/lib/binding/binding.rs @@ -49,6 +49,7 @@ use ruff_python_ast::name::Name; use ruff_text_size::Ranged; use ruff_text_size::TextRange; use starlark_map::small_set::SmallSet; +use thin_vec::ThinVec; use vec1::Vec1; use crate::alt::class::class_field::ClassField; @@ -926,7 +927,7 @@ pub enum Key { /// /// The resulting type may not actually involve a legacy type param, since it may turn out I am /// some other kind of type. - PossibleLegacyTParam(TextRange), + PossibleLegacyTParam(Box<(TextRange, TextRange)>), /// A `del` statement. It is a `Binding` associated with a the type `Any` because `del` defines a name in scope, /// so we need to provide a `Key` for any reads of that name in the edge case where there is no other definition /// @@ -965,7 +966,7 @@ impl Ranged for Key { Self::UsageLink(r) => *r, Self::YieldLink(r) => *r, Self::Delete(r) => *r, - Self::PossibleLegacyTParam(r) => *r, + Self::PossibleLegacyTParam(x) => x.0, Self::PatternNarrow(r) => *r, Self::Exhaustive(_, r) => *r, } @@ -1005,8 +1006,13 @@ impl DisplayWith for Key { Self::UsageLink(r) => write!(f, "Key::UsageLink({})", ctx.display(r)), Self::YieldLink(r) => write!(f, "Key::YieldLink({})", ctx.display(r)), Self::Delete(r) => write!(f, "Key::Delete({})", ctx.display(r)), - Self::PossibleLegacyTParam(r) => { - write!(f, "Key::PossibleLegacyTParam({})", ctx.display(r)) + Self::PossibleLegacyTParam(x) => { + write!( + f, + "Key::PossibleLegacyTParam({}, scope={})", + ctx.display(&x.0), + ctx.display(&x.1) + ) } Self::PatternNarrow(r) => write!(f, "Key::PatternNarrow({})", ctx.display(r)), Self::Exhaustive(kind, r) => { @@ -2337,7 +2343,7 @@ pub enum Binding { /// The TextRange is optional and controls whether to produce an error /// saying there are scoped type parameters for this function / class, and /// therefore the use of legacy type parameters is invalid. - PossibleLegacyTParam(Idx, Option), + PossibleLegacyTParam(ThinVec>, Option), /// An assignment to a name. NameAssign(Box), /// A type alias (legacy, scoped, or `TypeAliasType` call). @@ -2525,8 +2531,15 @@ impl DisplayWith for Binding { Self::TypeParameter(tp) => { write!(f, "TypeParameter({}, {}, ..)", tp.identity, tp.kind) } - Self::PossibleLegacyTParam(k, _) => { - write!(f, "PossibleLegacyTParam({})", ctx.display(*k)) + Self::PossibleLegacyTParam(keys, _) => { + write!(f, "PossibleLegacyTParam([")?; + for (i, &k) in keys.iter().enumerate() { + if i > 0 { + write!(f, ", ")?; + } + write!(f, "{}", ctx.display(k))?; + } + write!(f, "])") } Self::AnnotatedType(k1, k2) => { write!( diff --git a/pyrefly/lib/binding/bindings.rs b/pyrefly/lib/binding/bindings.rs index d7bcc391f6..2f8beabeb0 100644 --- a/pyrefly/lib/binding/bindings.rs +++ b/pyrefly/lib/binding/bindings.rs @@ -1799,10 +1799,11 @@ impl<'a> BindingsBuilder<'a> { // intercept_lookup which wraps them in PossibleLegacyTParam. // By finalize time we can follow through to the original // binding to check whether it's actually a type alias. - Binding::PossibleLegacyTParam(tparam_idx, _) - if let Some(legacy_binding) = self.idx_to_binding(*tparam_idx) => - { - self.follow_to_type_alias(legacy_binding.idx()) + Binding::PossibleLegacyTParam(tparam_indices, _) => { + tparam_indices.iter().find_map(|tparam_idx| { + let legacy_binding = self.idx_to_binding(*tparam_idx)?; + self.follow_to_type_alias(legacy_binding.idx()) + }) } _ => None, } @@ -2166,8 +2167,8 @@ impl LegacyTParamId { /// Note that the range here is not the range of the full `LegacyTParamId`, but /// just of the name being bound (which in the `Attr` case is just the base /// rather than the entire identifier). - fn as_possible_legacy_tparam_key(&self) -> Key { - Key::PossibleLegacyTParam(self.as_identifier().range) + fn as_possible_legacy_tparam_key(&self, scope_range: TextRange) -> Key { + Key::PossibleLegacyTParam(Box::new((self.as_identifier().range, scope_range))) } /// Get the key used to track this potential legacy tparam in the `legacy_tparams` map. @@ -2363,18 +2364,21 @@ impl<'a> BindingsBuilder<'a> { // Short circuit if there are too many forwards - it may mean there's a cycle. let mut original_binding = self.idx_to_binding(original_idx); let mut gas = Gas::new(100); - while let Some( - Binding::Forward(fwd_idx) - | Binding::PromoteForward(fwd_idx) - | Binding::ForwardToFirstUse(fwd_idx) - | Binding::Phi(JoinStyle::NarrowOf(fwd_idx), _), - ) = original_binding - { + loop { if gas.stop() { return None; - } else { - original_idx = *fwd_idx; - original_binding = self.idx_to_binding(original_idx); + } + match original_binding { + Some( + Binding::Forward(fwd_idx) + | Binding::PromoteForward(fwd_idx) + | Binding::ForwardToFirstUse(fwd_idx) + | Binding::Phi(JoinStyle::NarrowOf(fwd_idx), _), + ) => { + original_idx = *fwd_idx; + original_binding = self.idx_to_binding(original_idx); + } + _ => break, } } Some((original_idx, original_binding)) @@ -2392,22 +2396,62 @@ impl<'a> BindingsBuilder<'a> { has_scoped_type_params: bool, ) -> Option { let (original_idx, original_binding) = self.get_original_binding(original_idx)?; - // If we found a potential legacy type variable, first insert the key / binding pair - // for the raw lookup, then insert another key / binding pair for the - // `CheckLegacyTypeParam`, and return the `Idx`. let tparam_idx = Self::make_legacy_tparam(&id, original_binding, original_idx) .map(|(k, v)| self.insert_binding(k, v))?; - let idx = self.insert_binding( - id.as_possible_legacy_tparam_key(), - Binding::PossibleLegacyTParam( - tparam_idx, - if has_scoped_type_params { - Some(id.range()) - } else { - None - }, - ), - ); + + let scope_range = self.scopes.current_scope_range(); + let key = match &id { + LegacyTParamId::Name(_) => id.as_possible_legacy_tparam_key(scope_range), + LegacyTParamId::Attr(..) => { + let original_range = self.table.get::().0.idx_to_key(original_idx).range(); + Key::PossibleLegacyTParam(Box::new((original_range, scope_range))) + } + }; + let existing_idx = self.table.get::().0.key_to_idx(&key); + let idx = if let Some(existing_idx) = existing_idx { + let entry = self.table.get_mut::(); + if let Some(binding) = entry.1.get_mut(existing_idx) { + match binding { + Binding::PossibleLegacyTParam(tparams, _) => { + if !tparams.contains(&tparam_idx) { + tparams.push(tparam_idx); + } + } + _ => unreachable!("Expected Binding::PossibleLegacyTParam"), + } + existing_idx + } else { + let mut tparams = ThinVec::new(); + tparams.push(tparam_idx); + entry.1.insert( + existing_idx, + Binding::PossibleLegacyTParam( + tparams, + if has_scoped_type_params { + Some(id.range()) + } else { + None + }, + ), + ); + existing_idx + } + } else { + let mut tparams = ThinVec::new(); + tparams.push(tparam_idx); + self.insert_binding( + key, + Binding::PossibleLegacyTParam( + tparams, + if has_scoped_type_params { + Some(id.range()) + } else { + None + }, + ), + ) + }; + Some(PossibleTParam { id, idx, @@ -2472,8 +2516,12 @@ impl<'a> BindingsBuilder<'a> { for entry in legacy_tparams.legacy_tparams.values() { match entry { TParamLookupResult::MaybeTParam(possible_tparam) => { - self.scopes - .add_possible_legacy_tparam(possible_tparam.id.as_identifier()); + let key = self.table.get::().0.idx_to_key(possible_tparam.idx); + let range = key.range(); + self.scopes.add_possible_legacy_tparam( + possible_tparam.id.as_identifier().id.clone(), + range, + ); } _ => {} } diff --git a/pyrefly/lib/binding/scope.rs b/pyrefly/lib/binding/scope.rs index dd9441c1ce..bbec0d8e9b 100644 --- a/pyrefly/lib/binding/scope.rs +++ b/pyrefly/lib/binding/scope.rs @@ -279,11 +279,12 @@ impl MutableCapture { &self, name: &Name, kind: MutableCaptureKind, + scope_range: TextRange, ) -> Result { match &self.original { Result::Ok(static_info) => { if self.kind == kind { - Ok(static_info.as_key(name)) + Ok(static_info.as_key(name, scope_range)) } else { // TODO(stroxler): this error isn't quite right but preserves existing behavior Err(MutableCaptureError::AssignedBeforeNonlocal) @@ -355,7 +356,7 @@ impl StaticInfo { self.style.annotation() } - fn as_key(&self, name: &Name) -> Key { + fn as_key(&self, name: &Name, scope_range: TextRange) -> Key { let short_identifier = || { ShortIdentifier::new(&Identifier { node_index: AtomicNodeIndex::default(), @@ -372,7 +373,9 @@ impl StaticInfo { } StaticStyle::ImplicitGlobal => Key::ImplicitGlobal(Box::new(name.clone())), StaticStyle::SingleDef(..) => Key::Definition(short_identifier()), - StaticStyle::PossibleLegacyTParam => Key::PossibleLegacyTParam(self.range), + StaticStyle::PossibleLegacyTParam => { + Key::PossibleLegacyTParam(Box::new((self.range, scope_range))) + } } } @@ -1440,7 +1443,11 @@ impl ScopeTreeNode { } } for (name, info) in &self.scope.stat.0 { - if let Some(key) = table.types.0.key_to_idx(&info.as_key(name)) { + if let Some(key) = table + .types + .0 + .key_to_idx(&info.as_key(name, self.scope.range)) + { visitor(key); } } @@ -2565,12 +2572,12 @@ impl Scopes { /// of the scope, but only for static type lookups, and might potentially /// intercept the raw runtime value of a pre-PEP-695 legacy type variable /// to turn it into a quantified type parameter. - pub fn add_possible_legacy_tparam(&mut self, name: &Identifier) { + pub fn add_possible_legacy_tparam(&mut self, name: Name, range: TextRange) { self.current_mut().stat.upsert( - Hashed::new(name.id.clone()), - name.range, + Hashed::new(name), + range, StaticStyle::PossibleLegacyTParam, - name.range, + range, ) } @@ -3151,7 +3158,7 @@ impl Scopes { return None; } - let forward_ref_key = static_info.as_key(name.into_key()); + let forward_ref_key = static_info.as_key(name.into_key(), scope.range); return Some(NameReadInfo::Anywhere { key: forward_ref_key, // If we look up static info from the a non-barrier scope because we didn't find @@ -3262,7 +3269,8 @@ impl Scopes { style: StaticStyle::MutableCapture(capture), .. }) => { - let key = capture.key_or_error(name.into_key(), kind)?; + let key = + capture.key_or_error(name.into_key(), kind, self.current_scope_range())?; let implicit_builtin_module = capture .original .as_ref() @@ -3309,7 +3317,7 @@ impl ScopeTrace { Exportable::Initialized(*key, None) } } - None => Exportable::Uninitialized(static_info.as_key(name.into_key())), + None => Exportable::Uninitialized(static_info.as_key(name.into_key(), scope.range)), }; exportables.insert_hashed(name.owned(), exportable); } diff --git a/pyrefly/lib/state/ide.rs b/pyrefly/lib/state/ide.rs index c432088811..ceba08d6e8 100644 --- a/pyrefly/lib/state/ide.rs +++ b/pyrefly/lib/state/ide.rs @@ -109,8 +109,11 @@ fn find_definition_key_from<'a>(bindings: &'a Bindings, key: &'a Key) -> Option< Binding::Phi(_, branches) if !branches.is_empty() => { current_idx = branches[0].value_key } - Binding::PossibleLegacyTParam(k, _) => { - let binding = bindings.get(*k); + Binding::PossibleLegacyTParam(keys, _) => { + let first_key = keys + .first() + .expect("PossibleLegacyTParam bindings must always have at least one key"); + let binding = bindings.get(*first_key); current_idx = binding.idx(); } Binding::AssignToSubscript(x) @@ -146,8 +149,11 @@ fn create_intermediate_definition_from( Binding::Forward(k) | Binding::PromoteForward(k) | Binding::ForwardToFirstUse(k) => { current_binding = bindings.get(*k) } - Binding::PossibleLegacyTParam(k, _) => { - let binding = bindings.get(*k); + Binding::PossibleLegacyTParam(keys, _) => { + let first_key = keys + .first() + .expect("PossibleLegacyTParam bindings must always have at least one key"); + let binding = bindings.get(*first_key); current_binding = bindings.get(binding.idx()); } Binding::Import(x) => { diff --git a/pyrefly/lib/test/generic_legacy.rs b/pyrefly/lib/test/generic_legacy.rs index 696c002659..e1e440331c 100644 --- a/pyrefly/lib/test/generic_legacy.rs +++ b/pyrefly/lib/test/generic_legacy.rs @@ -933,7 +933,6 @@ class A: ); testcase!( - bug = "We currently create separate narrows for modules that may contain legacy type variables, we need to merge them", test_multiple_possible_legacy_tparams, TestEnv::one( "foo", @@ -943,22 +942,33 @@ testcase!( from typing import Generic, assert_type import foo -# Here, the `foo.C` possible-legacy-tparam binding is the one that winds up in scope, we -# lose track of the `foo.T` one. It probably doesn't matter very much since we at least -# understand the signature correctly. def f(x: foo.T, y: foo.C) -> foo.T: - z: foo.T = x # E: Type variable `T` is not in scope # E: `T` is not assignable to `TypeVar[T]` - return z # E: Returned type `TypeVar[T]` is not assignable to declared return type `T + z: foo.T = x + return z assert_type(f(1, foo.C()), int) -# The same thing happens here, but it's a much bigger problem because now we forget -# about the type variable identity for the entire class body, so the signatures come out -# wrong. class MyList(Generic[foo.T], list[tuple[foo.C, foo.T]]): - def my_append(self, c: foo.C, t: foo.T): # E: Type variable `T` is not in scope - self.append((c, t)) # E: Argument `tuple[C, TypeVar[T]]` is not assignable to parameter `object` with type `tuple[C, T]` in function `list.append` + def my_append(self, c: foo.C, t: foo.T): + self.append((c, t)) my_list: MyList[int] = MyList() -my_list.my_append(foo.C(), 5) # E: Argument `Literal[5]` is not assignable to parameter `t` with type `TypeVar[T]` in function `MyList.my_append` +my_list.my_append(foo.C(), 5) + "#, +); + +testcase!( + test_multiple_actual_legacy_tparams_from_same_module, + TestEnv::one( + "foo", + "from typing import TypeVar\nT = TypeVar('T')\nU = TypeVar('U')" + ), + r#" +from typing import assert_type +import foo + +def f(x: foo.T, y: foo.U) -> tuple[foo.T, foo.U]: + return (x, y) + +assert_type(f(1, "hello"), tuple[int, str]) "#, );