diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py index fe5eb1f3b..17e6f8916 100644 --- a/narwhals/_pandas_like/series.py +++ b/narwhals/_pandas_like/series.py @@ -18,6 +18,7 @@ from narwhals._pandas_like.utils import select_columns_by_name from narwhals._pandas_like.utils import set_axis from narwhals._pandas_like.utils import to_datetime +from narwhals.dependencies import is_numpy_scalar from narwhals.utils import Implementation from narwhals.utils import import_dtypes_module @@ -112,7 +113,7 @@ def __getitem__(self, idx: int) -> Any: ... def __getitem__(self, idx: slice | Sequence[int]) -> Self: ... def __getitem__(self, idx: int | slice | Sequence[int]) -> Any | Self: - if isinstance(idx, int): + if isinstance(idx, int) or is_numpy_scalar(idx): return self._native_series.iloc[idx] return self._from_native_series(self._native_series.iloc[idx]) diff --git a/narwhals/dependencies.py b/narwhals/dependencies.py index 463a64a67..b3d0dafdc 100644 --- a/narwhals/dependencies.py +++ b/narwhals/dependencies.py @@ -220,6 +220,11 @@ def is_numpy_array(arr: Any) -> TypeGuard[np.ndarray]: return (np := get_numpy()) is not None and isinstance(arr, np.ndarray) +def is_numpy_scalar(scalar: Any) -> TypeGuard[np.generic]: + """Check whether `scalar` is a NumPy Scalar without importing NumPy.""" + return (np := get_numpy()) is not None and np.isscalar(scalar) + + def is_pandas_like_dataframe(df: Any) -> bool: """Check whether `df` is a pandas-like DataFrame without doing any imports. diff --git a/narwhals/series.py b/narwhals/series.py index 2846aebea..d78ec945e 100644 --- a/narwhals/series.py +++ b/narwhals/series.py @@ -11,6 +11,7 @@ from typing import TypeVar from typing import overload +from narwhals.dependencies import is_numpy_scalar from narwhals.dtypes import _validate_dtype from narwhals.typing import IntoSeriesT from narwhals.utils import _validate_rolling_arguments @@ -67,7 +68,7 @@ def __getitem__(self: Self, idx: int) -> Any: ... def __getitem__(self: Self, idx: slice | Sequence[int]) -> Self: ... def __getitem__(self: Self, idx: int | slice | Sequence[int]) -> Any | Self: - if isinstance(idx, int): + if isinstance(idx, int) or is_numpy_scalar(idx): return self._compliant_series[idx] return self._from_compliant_series(self._compliant_series[idx]) diff --git a/tests/series_only/scalar_index_test.py b/tests/series_only/scalar_index_test.py new file mode 100644 index 000000000..21c9a10ba --- /dev/null +++ b/tests/series_only/scalar_index_test.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import narwhals.stable.v1 as nw + +np = nw.dependencies.get_numpy() +pd = nw.dependencies.get_pandas() + + +def test_index() -> None: + s = pd.Series([0, 1, 2]) + snw = nw.from_native(s, series_only=True) + assert snw[snw[0]] == np.int64(0)