Repository: numpy/numpy
File: numpy/lib/_shape_base_impl.pyi (Python)

Code before:
@type_check_only
class _ArrayWrap(Protocol):
    def __call__(
        self,
        array: NDArray[Any],
        context: tuple[ufunc, tuple[Any, ...], int] | None = ...,
        return_scalar: bool = ...,
        /,
    ) -> Any: ...

@type_check_only
class _SupportsArrayWrap(Protocol):
    @property
    def __array_wrap__(self) -> _ArrayWrap: ...

# Protocol for array-like objects that preserve their type through split operations.
# Requires shape (with __getitem__ for axis access), ndim for dimensional checks
# in hsplit/vsplit/dsplit, swapaxes for axis manipulation, and __getitem__ for
# slicing. Examples: pandas DataFrame, xarray DataArray.
@type_check_only
class _SupportsSplitOps(Protocol):
    @property
    def shape(self) -> tuple[int, ...]: ...
    @property
    def ndim(self) -> int: ...
    def swapaxes(self, axis1: SupportsIndex, axis2: SupportsIndex, /) -> Self: ...
    def __getitem__(self, key: Any, /) -> Self: ...

###

def take_along_axis[ScalarT: np.generic](
    arr: ScalarT | NDArray[ScalarT],
    indices: NDArray[integer],
    axis: int | None = -1,
) -> NDArray[ScalarT]: ...

def put_along_axis[ScalarT: np.generic](
    arr: NDArray[ScalarT],
    indices: NDArray[integer],
    values: ArrayLike,
    axis: int | None,
) -> None: ...

@overload
def apply_along_axis[**Tss, ScalarT: np.generic](
    func1d: Callable[Concatenate[np.ndarray, Tss], _ArrayLike[ScalarT]],
    axis: SupportsIndex,
    arr: ArrayLike,
    *args: Tss.args,
    **kwargs: Tss.kwargs,
) -> NDArray[ScalarT]: ...

Review comment (suggestion):
I'm not sure if this would cause any issues for pandas, but the `axis{1,2}` are in contravariant (input) positions, so it would be more general to annotate them as `: int`.

Code after:
@type_check_only
class _ArrayWrap(Protocol):
    def __call__(
        self,
        array: NDArray[Any],
        context: tuple[ufunc, tuple[Any, ...], int] | None = ...,
        return_scalar: bool = ...,
        /,
    ) -> Any: ...

@type_check_only
class _SupportsArrayWrap(Protocol):
    @property
    def __array_wrap__(self) -> _ArrayWrap: ...

# Protocol for array-like objects that preserve their type through split operations.
# Requires shape for size, ndim for dimensional checks in hsplit/vsplit/dsplit,
# swapaxes for axis manipulation, and __getitem__ for slicing.
@type_check_only
class _SupportsSplitOps(Protocol):
    @property
    def shape(self) -> tuple[int, ...]: ...
    @property
    def ndim(self) -> int: ...
    def swapaxes(self, axis1: int, axis2: int, /) -> Self: ...
    def __getitem__(self, key: Any, /) -> Self: ...

###

def take_along_axis[ScalarT: np.generic](
    arr: ScalarT | NDArray[ScalarT],
    indices: NDArray[integer],
    axis: int | None = -1,
) -> NDArray[ScalarT]: ...

def put_along_axis[ScalarT: np.generic](
    arr: NDArray[ScalarT],
    indices: NDArray[integer],
    values: ArrayLike,
    axis: int | None,
) -> None: ...

@overload
def apply_along_axis[**Tss, ScalarT: np.generic](
    func1d: Callable[Concatenate[np.ndarray, Tss], _ArrayLike[ScalarT]],
    axis: SupportsIndex,
    arr: ArrayLike,
    *args: Tss.args,
    **kwargs: Tss.kwargs,
) -> NDArray[ScalarT]: ...
@overload