Validated type-hint changes

Author: peterdsharpe

aerosandbox/numpy/arithmetic_dyadic.py

@@ -29,7 +29,7 @@ def _make_casadi_types_broadcastable(
     tuple[Array, Array]
         Both arrays tiled to have the same (broadcast) shape.
     """
-    def shape_2D(object: float | int | Iterable | _onp.ndarray) -> tuple:
+    def shape_2D(object: Vectorizable) -> tuple:
         shape = _onp.shape(object)
         if len(shape) == 0:
             return (1, 1)

aerosandbox/numpy/arithmetic_monadic.py

@@ -6,6 +6,7 @@
 import numpy as _onp
 import casadi as _cas
 from aerosandbox.numpy.determine_type import is_casadi_type
+from aerosandbox.numpy.array import asarray
 from aerosandbox.numpy.typing import Array, ArrayLike, Scalar
 
 
@@ -80,6 +81,7 @@ def mean(x: ArrayLike, axis: int | None = None) -> Scalar | Array:
         return _onp.mean(x, axis=axis)
 
     else:
+        x = asarray(x)  # Ensure x is Array for .shape access
         if axis == 0:
             return sum(x, axis=0) / x.shape[0]
         elif axis == 1:

aerosandbox/numpy/array.py

@@ -320,6 +320,7 @@ def length(array: ArrayLike | Scalar) -> int:
             return 1
 
     else:
+        array = asarray(array)  # Ensure array is Array for .shape access
         if array.shape[0] != 1:
             return array.shape[0]
         else:
@@ -359,6 +360,7 @@ def diag(v: ArrayLike, k: int = 0) -> Array:
         return _onp.diag(v, k=k)
 
     else:
+        v = asarray(v)  # Ensure v is Array for .shape access
         if 1 in v.shape:  # If v is a 1D array, construct a diagonal matrix
             if v.shape[0] == 1:
                 v = v.T
@@ -421,6 +423,7 @@ def roll(a: ArrayLike, shift: int | tuple[int, ...], axis: int | tuple[int, ...]
     if not is_casadi_type(a, recursive=False):
         return _onp.roll(a, shift, axis=axis)
     else:
+        a = asarray(a)  # Ensure a is Array for .shape access
         if axis is None:
             a_flat = reshape(a, -1)
             result = roll(a_flat, shift, axis=0)
@@ -480,6 +483,7 @@ def max(a: ArrayLike, axis: int | None = None) -> Scalar | Array:
         )
 
     else:
+        a = asarray(a)  # Ensure a is Array for .shape access
         if axis is None:
             return _cas.mmax(a)
 
@@ -532,6 +536,7 @@ def min(a: ArrayLike, axis: int | None = None) -> Scalar | Array:
         )
 
     else:
+        a = asarray(a)  # Ensure a is Array for .shape access
         if axis is None:
             return _cas.mmin(a)
 

aerosandbox/numpy/calculus.py

@@ -290,13 +290,15 @@ def trapz(x: ArrayLike, modify_endpoints: bool = False) -> Array:
     integrate_discrete_intervals : Preferred alternative.
     """
     import warnings
+    from aerosandbox.numpy.array import asarray
 
     warnings.warn(
         "trapz() will eventually be deprecated, since NumPy plans to remove it in the upcoming NumPy 2.0 release (2024). \n"
         'For discrete intervals, use asb.numpy.integrate_discrete_intervals(f, method="trapz") instead.',
         PendingDeprecationWarning,
     )
 
+    x = asarray(x)  # Convert to Array for subscripting
     integral = (x[1:] + x[:-1]) / 2
     if modify_endpoints:
         integral[0] = integral[0] + x[0] * 0.5

aerosandbox/numpy/conditionals.py

@@ -10,15 +10,15 @@
 
 
 def where(
-    condition: Array,
+    condition: Array | bool | _onp.bool_,
     value_if_true: Vectorizable,
     value_if_false: Vectorizable,
 ) -> Array:
     """Return elements chosen from x or y depending on condition.
 
     Parameters
     ----------
-    condition : Array
+    condition : Array | bool
         Where True, yield ``value_if_true``, otherwise yield ``value_if_false``.
     value_if_true : Vectorizable
         Values from which to choose where ``condition`` is True.

aerosandbox/numpy/integrate.py

@@ -251,6 +251,7 @@ def solve_ivp(
         * `t_variable` is a CasADi variable (cas.MX)
         * `y_variables` is a CasADi variable (cas.MX), possibly a vector of variables
         """
+        assert y_variables is not None  # Type narrowing for type checker
 
         t0 = t_span[0]
         tf = t_span[1]

aerosandbox/numpy/integrate_discrete.py

@@ -1,12 +1,13 @@
 from typing import Literal
 import casadi as _cas
 import numpy as _onp
-from aerosandbox.numpy.array import length, concatenate
+from aerosandbox.numpy.array import length, concatenate, asarray
+from aerosandbox.numpy.typing import ArrayLike
 
 
 def integrate_discrete_intervals(
-    f: _onp.ndarray | _cas.MX,
-    x: _onp.ndarray | _cas.MX | None = None,
+    f: ArrayLike,
+    x: ArrayLike | None = None,
     multiply_by_dx: bool = True,
     method: Literal[
         "forward_euler",
@@ -50,10 +51,15 @@ def integrate_discrete_intervals(
             - "periodic"
 
     """
+    # Convert inputs to arrays for subscripting
+    f = asarray(f)
+
     # Determine if an x-array was specified, and calculate dx.
     x_is_specified = x is not None
     if not x_is_specified:
         x = _onp.arange(length(f))
+    else:
+        x = asarray(x)
 
     dx = x[1:] - x[:-1]
 
@@ -266,8 +272,8 @@ def integrate_discrete_intervals(
 
 
 def integrate_discrete_squared_curvature(
-    f: _onp.ndarray | _cas.MX,
-    x: _onp.ndarray | _cas.MX | None = None,
+    f: ArrayLike,
+    x: ArrayLike | None = None,
     method: Literal[
         "cubic", "simpson", "hybrid_simpson_cubic"
     ] = "hybrid_simpson_cubic",
@@ -325,10 +331,15 @@ def integrate_discrete_squared_curvature(
             well as a regularization strategy. (It is still convergent to the true value in the high-sample-rate limit.)
 
     """
+    # Convert inputs to arrays for subscripting
+    f = asarray(f)
+
     # Determine if an x-array was specified, and calculate dx.
     x_is_specified = x is not None
     if not x_is_specified:
         x = _onp.arange(length(f))
+    else:
+        x = asarray(x)
 
     if method in ["cubic", "cubic_spline"]:
         x1 = x[:-3]

aerosandbox/numpy/linalg.py

@@ -7,6 +7,7 @@
 import casadi as _cas
 from aerosandbox.numpy.arithmetic_monadic import sum, abs
 from aerosandbox.numpy.determine_type import is_casadi_type
+from aerosandbox.numpy.array import asarray
 from aerosandbox.numpy.typing import ArrayLike, Array, Scalar, VectorLike
 from numpy.linalg import *
 
@@ -71,6 +72,7 @@ def outer(x: VectorLike, y: VectorLike, manual: bool = False) -> Array:
         return _onp.outer(x, y)
 
     else:
+        y = asarray(y)  # Ensure y is Array for .shape access
         if len(y.shape) == 1:  # Force y to be transposable if it's not.
             y = _onp.expand_dims(y, 1)
         return x @ y.T
@@ -214,6 +216,7 @@ def norm(
         return _onp.linalg.norm(x, ord=ord, axis=axis, keepdims=keepdims)
 
     else:
+        x = asarray(x)  # Ensure x is Array for .shape access
         # Figure out which axis, if any, to take a vector norm about.
         if axis is not None:
             if not (axis == 0 or axis == 1 or axis == -1):

aerosandbox/numpy/typing.py

@@ -122,16 +122,18 @@
 # INPUT TYPES - Permissive types for function parameters (hybrid)
 # =============================================================================
 
-VectorLike = Sequence[float] | Sequence[int] | _onp.ndarray | CasADiType
-"""Permissive input for vector parameters: sequence, ndarray, or CasADi.
+VectorLike = int | float | _onp.integer | _onp.floating | Sequence[int | float] | _onp.ndarray | CasADiType
+"""Permissive input for vector parameters: scalar, sequence, ndarray, or CasADi.
 
 Use for function INPUTS that will be converted to Vector internally.
+Includes scalars so that `Scalar | Array` return values can be passed directly.
 """
 
-ArrayLike = Sequence[float] | Sequence[int] | _onp.ndarray | CasADiType
-"""Permissive input for array parameters: sequence, ndarray, or CasADi.
+ArrayLike = int | float | _onp.integer | _onp.floating | Sequence[int | float] | _onp.ndarray | CasADiType
+"""Permissive input for array parameters: scalar, sequence, ndarray, or CasADi.
 
 Use for function INPUTS that will be converted to Array internally.
+Includes scalars so that `Scalar | Array` return values can be passed directly.
 """
 
 PointLike = Sequence[float] | _onp.ndarray | CasADiType