diff --git a/crates/pyrefly_config/src/error_kind.rs b/crates/pyrefly_config/src/error_kind.rs index 719edff13a..407455cfc3 100644 --- a/crates/pyrefly_config/src/error_kind.rs +++ b/crates/pyrefly_config/src/error_kind.rs @@ -407,6 +407,8 @@ pub enum ErrorKind { UnannotatedReturn, /// Attempting to use a name that may be unbound or uninitialized UnboundName, + /// A type parameter may not be constrained by some valid calls to a generic function. + UnconstrainedTypeVar, /// An error caused by a keyword argument used in the wrong place. UnexpectedKeyword, /// An error caused by passing a positional argument for a keyword-only parameter. diff --git a/pyrefly/lib/alt/expr.rs b/pyrefly/lib/alt/expr.rs index b8f874ed60..d84fb0c961 100644 --- a/pyrefly/lib/alt/expr.rs +++ b/pyrefly/lib/alt/expr.rs @@ -2767,7 +2767,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let construct = kind.to_string(); let mut arg_name = false; let mut restriction = None; - let mut default = None; + let mut default: Option<&Expr> = None; let mut variance = None; let check_name_arg = |arg: &Expr| { @@ -2846,14 +2846,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } } "default" => { - default = Some(( - self.expr_untype( - &kw.value, - TypeFormContext::quantified_kind_default(kind), - errors, - ), - kw.value.range(), - )) + default = Some(&kw.value); } "covariant" => try_set_variance(kw, PreInferenceVariance::Covariant), "contravariant" => try_set_variance(kw, PreInferenceVariance::Contravariant), @@ -2918,12 +2911,23 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } let restriction = restriction.unwrap_or(Restriction::Unrestricted); let mut default_value = None; - if let Some((default_ty, default_range)) = default { + if let Some(default_expr) = default { + let default_ty = if let Some(default) = + self.parse_int_tuple_type_var_default(default_expr, &restriction, errors) + { + default + } else { + self.expr_untype( + default_expr, + TypeFormContext::quantified_kind_default(kind), + errors, + ) + }; default_value = Some(self.validate_type_var_default( &name.id, kind, &default_ty, - default_range, + default_expr.range(), &restriction, errors, )); @@ -4985,6 +4989,31 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { .map(IntTuple::from_types) } + /// Parse a list-form default for an `IntTuple`-bounded type variable. + pub(crate) fn parse_int_tuple_type_var_default( + &self, + default: &Expr, + restriction: &Restriction, + errors: &ErrorCollector, + ) -> Option { + let Restriction::Bound(bound) = restriction else { + return None; + }; + let Expr::List(list) = default else { + return None; + }; + if !self.solver().config.tensor_shapes + || !is_int_tuple_bound(bound, &self.stdlib.int().clone().to_type()) + { + return None; + } + Some( + self.parse_int_tuple_shape_args(&list.elts, TypeFormContext::TypeVarDefault, errors) + .map(|shape| shape.to_shape_arg_type()) + .unwrap_or_else(Type::any_error), + ) + } + /// Parse a registered shaped-array annotation. /// /// The registered shape parameter is a single ordinary type argument that diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index a61327b261..21a1eabc4a 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -906,7 +906,13 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let tparams = self.collect_jaxtyping_tparams(&callable, &def.tparams, stmt.name.range, errors); - self.validate_shape_extension_function_parameters(stmt, &def.params, &tparams, errors); + self.validate_shape_extension_function_parameters( + stmt, + &def.params, + &callable.ret, + &tparams, + errors, + ); let mut metadata = def.metadata.clone(); self.record_shape_flag_constructor_sources( diff --git a/pyrefly/lib/alt/shape_extension.rs b/pyrefly/lib/alt/shape_extension.rs index 6a89dfe37a..469cac9786 100644 --- a/pyrefly/lib/alt/shape_extension.rs +++ b/pyrefly/lib/alt/shape_extension.rs @@ -93,6 +93,48 @@ pub(crate) fn direct_function_parameter_sources( .collect() } +#[derive(Clone, Copy, PartialEq, Eq)] +enum TypeParameterCoverage { + None, + Partial, + Total, +} + +fn type_mentions_parameter(ty: &Type, tparam: &Quantified) -> bool { + let mut found = false; + ty.for_each_quantified(&mut |candidate| found |= candidate == tparam); + found +} + +/// Classify whether every top-level union arm mentions `tparam`. +/// +/// Resolved signatures have already expanded ordinary aliases. Other nested type structure is +/// intentionally treated as atomic because this shape-specific check does not model general +/// generic inference. +fn type_parameter_coverage(ty: &Type, tparam: &Quantified) -> TypeParameterCoverage { + match ty { + Type::Union(union) => { + if union + .members + .iter() + .all(|member| type_mentions_parameter(member, tparam)) + { + TypeParameterCoverage::Total + } else if union + .members + .iter() + .any(|member| type_mentions_parameter(member, tparam)) + { + TypeParameterCoverage::Partial + } else { + TypeParameterCoverage::None + } + } + _ if type_mentions_parameter(ty, tparam) => TypeParameterCoverage::Total, + _ => TypeParameterCoverage::None, + } +} + impl AnswersSolver<'_, '_, Ans> { pub(crate) fn validate_shape_extension_type_parameter_default( &self, @@ -118,11 +160,74 @@ impl AnswersSolver<'_, '_, Ans> { &self, stmt: &FunctionDefData, params: &[Param], + ret: &Type, tparams: &TParams, errors: &ErrorCollector, ) { self.validate_shape_flag_function_parameters(stmt, params, tparams, errors); self.validate_shape_index_function_parameters(stmt, params, tparams, errors); + self.validate_int_tuple_function_parameters(stmt, params, ret, tparams, errors); + } + + fn validate_int_tuple_function_parameters( + &self, + stmt: &FunctionDefData, + params: &[Param], + ret: &Type, + tparams: &TParams, + errors: &ErrorCollector, + ) { + if !self.solver().config.tensor_shapes { + return; + } + // `IntTuple` retains distinct provenance. A plain `tuple[int, ...]` bound may instead + // belong to ordinary typing or an implicit jaxtyping parameter, neither of which can use + // this shape-specific default. + for tparam in tparams.iter().filter(|tparam| { + tparam.default().is_none() + && matches!(tparam.restriction(), Restriction::Bound(Type::IntTuple(_))) + }) { + let mut appears_in_return = false; + ret.for_each_quantified(&mut |candidate| appears_in_return |= candidate == tparam); + if !appears_in_return { + continue; + } + let (has_required_source, has_partial_source) = params.iter().fold( + (false, false), + |(has_required_source, has_partial_source), param| { + let coverage = type_parameter_coverage(param.as_type(), tparam); + ( + has_required_source + || param.is_required() && coverage == TypeParameterCoverage::Total, + has_partial_source || coverage == TypeParameterCoverage::Partial, + ) + }, + ); + // This is intentionally a narrow lint for partially constraining unions. Optional + // and variadic sources can also be omitted, but overloads and runtime arity rules + // make broader definition-site analysis too noisy. + if has_partial_source && !has_required_source { + let range = stmt + .type_params + .as_ref() + .and_then(|params| { + params + .type_params + .iter() + .find(|param| param.name().id == *tparam.name()) + }) + .map_or_else(|| stmt.name.range(), Ranged::range); + self.error( + errors, + range, + ErrorKind::UnconstrainedTypeVar, + format!( + "`IntTuple` type parameter `{}` may be unconstrained for some calls; give it a default", + tparam.name(), + ), + ); + } + } } pub(crate) fn reject_legacy_shape_extension_bound( diff --git a/pyrefly/lib/alt/solve.rs b/pyrefly/lib/alt/solve.rs index 46b337c5a5..ada51d262d 100644 --- a/pyrefly/lib/alt/solve.rs +++ b/pyrefly/lib/alt/solve.rs @@ -2126,6 +2126,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { && let Some(n) = i.as_i64() { Type::Int(Int::Literal(n)) + } else if let Some(default) = + self.parse_int_tuple_type_var_default(default_expr, &restriction, errors) + { + default } else { self.expr_untype( default_expr, diff --git a/pyrefly/lib/test/shape_dsl.rs b/pyrefly/lib/test/shape_dsl.rs index 4421af113a..656541876e 100644 --- a/pyrefly/lib/test/shape_dsl.rs +++ b/pyrefly/lib/test/shape_dsl.rs @@ -6324,6 +6324,86 @@ def check(default: Array, explicit: Array[IntTuple]) -> None: "#, ); +testcase!( + test_empty_int_tuple_defaults_for_array_like_unions, + shape_extensions_env(), + r#" +from typing import assert_type +from shape_extensions import IntTuple +from typing_extensions import TypeVar + +class Array[Shape: IntTuple]: ... +class ndarray[Shape: IntTuple]: ... + +type ArrayLike[Shape: IntTuple = []] = ndarray[Shape] | Array[Shape] | float + +LegacyShape = TypeVar("LegacyShape", bound=IntTuple, default=[]) +LegacyMissing = TypeVar("LegacyMissing", bound=IntTuple) + +def direct[Shape: IntTuple = []]( + value: ndarray[Shape] | Array[Shape] | float, +) -> Array[Shape]: ... + +def through_alias[Shape: IntTuple = []](value: ArrayLike[Shape]) -> Array[Shape]: ... + +def concrete_default[Shape: IntTuple = [2, 3]]() -> Array[Shape]: ... + +def alias_without_function_default[Shape: IntTuple]( # E: `IntTuple` type parameter `Shape` may be unconstrained for some calls; give it a default + value: ArrayLike[Shape], +) -> Array[Shape]: ... + +def missing_default_direct[Shape: IntTuple]( # E: `IntTuple` type parameter `Shape` may be unconstrained for some calls; give it a default + value: ndarray[Shape] | Array[Shape] | float, +) -> Array[Shape]: ... + +type NestedArrayLike[Shape: IntTuple] = ArrayLike[Shape] + +def missing_default_through_alias[Shape: IntTuple]( # E: `IntTuple` type parameter `Shape` may be unconstrained for some calls; give it a default + value: NestedArrayLike[Shape], +) -> Array[Shape]: ... + +def suppressed[Shape: IntTuple]( # pyrefly: ignore[unconstrained-type-var] + value: ArrayLike[Shape], +) -> Array[Shape]: ... + +# Optional and variadic sources are intentionally outside this lint's ArrayLike-union scope. +def optional_only[Shape: IntTuple](value: Array[Shape] = ...) -> Array[Shape]: ... + +def variadic_only[Shape: IntTuple](*values: Array[Shape]) -> Array[Shape]: ... + +def unpacked_variadic[Shape: IntTuple](*values: *Shape) -> Array[Shape]: ... + +def bound_by_required_parameter[Shape: IntTuple]( + value: ArrayLike[Shape], required: Array[Shape], +) -> Array[Shape]: ... + +def unobservable[Shape: IntTuple](value: ArrayLike[Shape]) -> None: ... + +def legacy(value: ndarray[LegacyShape] | Array[LegacyShape] | float) -> Array[LegacyShape]: ... + +def legacy_missing_default( # E: `IntTuple` type parameter `LegacyMissing` may be unconstrained for some calls; give it a default + value: ndarray[LegacyMissing] | Array[LegacyMissing] | float, +) -> Array[LegacyMissing]: ... + +def bare_alias(value: ArrayLike) -> ArrayLike: ... + +# Function type parameter defaults provide the useful scalar fallback. +assert_type(direct(1.0), Array[[]]) +assert_type(through_alias(1.0), Array[[]]) +assert_type(legacy(1.0), Array[[]]) +assert_type(concrete_default(), Array[[2, 3]]) + +# The alias default only specializes a bare alias to ArrayLike[[]]. +assert_type(bare_alias(1.0), ArrayLike[[]]) + +def preserve_shape(array: Array[[2, 3]], nd: ndarray[[4]]) -> None: + assert_type(direct(array), Array[[2, 3]]) + assert_type(through_alias(nd), Array[[4]]) + assert_type(legacy(array), Array[[2, 3]]) + bare_alias(array) # E: is not assignable to parameter `value` +"#, +); + testcase!( test_tensor_shapes_gradual_size, legacy_shaped_array_env(), @@ -7874,6 +7954,10 @@ def concrete(x: Float[Array, "3 4"]) -> None: def named_variadic(x: Float[Array, "*batch channels"]) -> None: reveal_type(x) # E: revealed type: Shaped[Array, "*batch channels"] + +def scalar_or_variadic( + x: Float[Array, "*batch"] | float, +) -> Float[Array, "*batch"]: ... "#, ); @@ -14337,6 +14421,8 @@ T_co = TypeVar("T_co", bound=tuple[int, ...], covariant=True) class InvariantBox(Generic[T]): ... class CovariantBox(Generic[T_co]): ... +def scalar_or_box(value: InvariantBox[T] | float) -> InvariantBox[T]: ... + def check( invariant_concrete: InvariantBox[tuple[int, int]], covariant_concrete: CovariantBox[tuple[int, int]], diff --git a/scripts/error_presets.json b/scripts/error_presets.json index 7a8d0d25d4..0aa53863ba 100644 --- a/scripts/error_presets.json +++ b/scripts/error_presets.json @@ -122,6 +122,7 @@ "unannotated-protocol-member": ["legacy", "default", "strict", "all"], "unannotated-return": ["all"], "unbound-name": ["default", "strict", "all"], + "unconstrained-type-var": ["legacy", "default", "strict", "all"], "unexpected-keyword": ["basic", "legacy", "default", "strict", "all"], "unexpected-positional-argument": ["basic", "legacy", "default", "strict", "all"], "unimported-directive": ["legacy", "default", "strict", "all"], diff --git a/website/docs/error-kinds.mdx b/website/docs/error-kinds.mdx index 91cc775ae1..2f4ae23ed2 100644 --- a/website/docs/error-kinds.mdx +++ b/website/docs/error-kinds.mdx @@ -1964,6 +1964,32 @@ def f(check: bool): Compare this with [unknown-name](#unknown-name), which is reported when no definition at all is found for a name. +## unconstrained-type-var + +This is a type-aware lint to make errors writing Pyrefly tensor shape stubs +less likely. Pyrefly reports it when an `IntTuple`-bounded type parameter is +used in the return type, has no default, and appears in only some arms of a +top-level union in a function parameter (after expanding type alias). In that +situation, some calls will fail to constrain the parameter which might lead to +unexpended gradual shapes. + +Its main purpose is to catch a missing default in `ArrayLike`-style APIs, where +a scalar argument is meant to imply an empty shape. In that situtation, the desired +behavior requires setting the empty shape as a default: + +```python +from shape_extensions import IntTuple + +class Array[Shape: IntTuple]: ... +type ArrayLike[Shape: IntTuple] = Array[Shape] | float + +def as_array[Shape: IntTuple]( # error: `IntTuple` type parameter `Shape` may be unconstrained for some calls; give it a default + value: ArrayLike[Shape], +) -> Array[Shape]: ... + +def fixed[Shape: IntTuple = []](value: ArrayLike[Shape]) -> Array[Shape]: ... +``` + ## unexpected-keyword A function was called with an extra keyword argument.