Add validation that Shape is always solved [rfc] (#4898)

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.

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<Type> {
+        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

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(

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<Ans: LookupAnswer> AnswersSolver<'_, '_, Ans> {
     pub(crate) fn validate_shape_extension_type_parameter_default(
         &self,
@@ -118,11 +160,74 @@ impl<Ans: LookupAnswer> 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(

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,

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]],

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"],

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.