2025-summer/team-10
10
Fork 0
Code Issues Pull requests Projects Releases Activity
team-10/env/Lib/site-packages/narwhals/_utils.py

2024 lines
64 KiB
Python
Raw Normal View History

integrata generazione immagini
2025-08-02 07:34:44 +02:00
from __future__ import annotations
import os
import re
from collections.abc import Collection, Container, Iterable, Iterator, Mapping, Sequence
from datetime import timezone
from enum import Enum, auto
from functools import cache, lru_cache, wraps
from importlib.util import find_spec
from inspect import getattr_static, getdoc
from secrets import token_hex
from typing import (
TYPE_CHECKING,
Any,
Callable,
Generic,
Literal,
Protocol,
TypeVar,
Union,
cast,
overload,
)
from narwhals._enum import NoAutoEnum
from narwhals._exceptions import issue_deprecation_warning
from narwhals._typing_compat import assert_never, deprecated
from narwhals.dependencies import (
get_cudf,
get_dask_dataframe,
get_duckdb,
get_ibis,
get_modin,
get_pandas,
get_polars,
get_pyarrow,
get_pyspark_connect,
get_pyspark_sql,
get_sqlframe,
is_narwhals_series,
is_narwhals_series_int,
is_numpy_array_1d,
is_numpy_array_1d_int,
is_pandas_like_dataframe,
is_pandas_like_series,
is_polars_series,
is_pyarrow_chunked_array,
)
from narwhals.exceptions import ColumnNotFoundError, DuplicateError, InvalidOperationError
if TYPE_CHECKING:
from collections.abc import Set # noqa: PYI025
from types import ModuleType
import pandas as pd
import polars as pl
import pyarrow as pa
from typing_extensions import (
Concatenate,
LiteralString,
ParamSpec,
Self,
TypeAlias,
TypeIs,
)
from narwhals._compliant import (
CompliantExpr,
CompliantExprT,
CompliantFrameT,
CompliantSeriesOrNativeExprT_co,
CompliantSeriesT,
NativeFrameT_co,
NativeSeriesT_co,
)
from narwhals._compliant.typing import EvalNames
from narwhals._namespace import EagerAllowedImplementation, Namespace
from narwhals._translate import ArrowStreamExportable, IntoArrowTable, ToNarwhalsT_co
from narwhals.dataframe import DataFrame, LazyFrame
from narwhals.dtypes import DType
from narwhals.series import Series
from narwhals.typing import (
CompliantDataFrame,
CompliantLazyFrame,
CompliantSeries,
DataFrameLike,
DTypes,
IntoSeriesT,
MultiIndexSelector,
SingleIndexSelector,
SizedMultiIndexSelector,
SizeUnit,
SupportsNativeNamespace,
TimeUnit,
_1DArray,
_SliceIndex,
_SliceName,
_SliceNone,
)
FrameOrSeriesT = TypeVar(
"FrameOrSeriesT", bound=Union[LazyFrame[Any], DataFrame[Any], Series[Any]]
)
_T1 = TypeVar("_T1")
_T2 = TypeVar("_T2")
_T3 = TypeVar("_T3")
_T4 = TypeVar("_T4")
_T5 = TypeVar("_T5")
_T6 = TypeVar("_T6")
_T7 = TypeVar("_T7")
_Fn = TypeVar("_Fn", bound="Callable[..., Any]")
P = ParamSpec("P")
R = TypeVar("R")
R1 = TypeVar("R1")
R2 = TypeVar("R2")
class _SupportsVersion(Protocol):
__version__: str
class _SupportsGet(Protocol): # noqa: PYI046
def __get__(self, instance: Any, owner: Any | None = None, /) -> Any: ...
class _StoresColumns(Protocol):
@property
def columns(self) -> Sequence[str]: ...
_T = TypeVar("_T")
NativeT_co = TypeVar("NativeT_co", covariant=True)
CompliantT_co = TypeVar("CompliantT_co", covariant=True)
_ContextT = TypeVar("_ContextT", bound="_FullContext")
_Method: TypeAlias = "Callable[Concatenate[_ContextT, P], R]"
_Constructor: TypeAlias = "Callable[Concatenate[_T, P], R2]"
class _StoresNative(Protocol[NativeT_co]): # noqa: PYI046
"""Provides access to a native object.
Native objects have types like:
>>> from pandas import Series
>>> from pyarrow import Table
"""
@property
def native(self) -> NativeT_co:
"""Return the native object."""
...
class _StoresCompliant(Protocol[CompliantT_co]): # noqa: PYI046
"""Provides access to a compliant object.
Compliant objects have types like:
>>> from narwhals._pandas_like.series import PandasLikeSeries
>>> from narwhals._arrow.dataframe import ArrowDataFrame
"""
@property
def compliant(self) -> CompliantT_co:
"""Return the compliant object."""
...
class _StoresBackendVersion(Protocol):
@property
def _backend_version(self) -> tuple[int, ...]:
"""Version tuple for a native package."""
...
class _StoresVersion(Protocol):
_version: Version
"""Narwhals API version (V1 or MAIN)."""
class _StoresImplementation(Protocol):
_implementation: Implementation
"""Implementation of native object (pandas, Polars, PyArrow, ...)."""
class _LimitedContext(_StoresImplementation, _StoresVersion, Protocol):
"""Provides 2 attributes.
- `_implementation`
- `_version`
"""
class _FullContext(_StoresBackendVersion, _LimitedContext, Protocol):
"""Provides 3 attributes.
- `_implementation`
- `_backend_version`
- `_version`
"""
class ValidateBackendVersion(_StoresImplementation, Protocol):
"""Ensure the target `Implementation` is on a supported version."""
def _validate_backend_version(self) -> None:
"""Raise if installed version below `nw._utils.MIN_VERSIONS`.
**Only use this when moving between backends.**
Otherwise, the validation will have taken place already.
"""
_ = self._implementation._backend_version()
class Version(Enum):
V1 = auto()
V2 = auto()
MAIN = auto()
@property
def namespace(self) -> type[Namespace[Any]]:
if self is Version.V1:
from narwhals.stable.v1._namespace import Namespace as NamespaceV1
return NamespaceV1
if self is Version.V2:
from narwhals.stable.v2._namespace import Namespace as NamespaceV2
return NamespaceV2
from narwhals._namespace import Namespace
return Namespace
@property
def dtypes(self) -> DTypes:
if self is Version.V1:
from narwhals.stable.v1 import dtypes as dtypes_v1
return dtypes_v1
if self is Version.V2:
from narwhals.stable.v2 import dtypes as dtypes_v2
return dtypes_v2
from narwhals import dtypes
return dtypes
@property
def dataframe(self) -> type[DataFrame[Any]]:
if self is Version.V1:
from narwhals.stable.v1 import DataFrame as DataFrameV1
return DataFrameV1
if self is Version.V2:
from narwhals.stable.v2 import DataFrame as DataFrameV2
return DataFrameV2
from narwhals.dataframe import DataFrame
return DataFrame
@property
def lazyframe(self) -> type[LazyFrame[Any]]:
if self is Version.V1:
from narwhals.stable.v1 import LazyFrame as LazyFrameV1
return LazyFrameV1
if self is Version.V2:
from narwhals.stable.v2 import LazyFrame as LazyFrameV2
return LazyFrameV2
from narwhals.dataframe import LazyFrame
return LazyFrame
@property
def series(self) -> type[Series[Any]]:
if self is Version.V1:
from narwhals.stable.v1 import Series as SeriesV1
return SeriesV1
if self is Version.V2:
from narwhals.stable.v2 import Series as SeriesV2
return SeriesV2
from narwhals.series import Series
return Series
class Implementation(NoAutoEnum):
"""Implementation of native object (pandas, Polars, PyArrow, ...)."""
PANDAS = "pandas"
"""pandas implementation."""
MODIN = "modin"
"""Modin implementation."""
CUDF = "cudf"
"""cuDF implementation."""
PYARROW = "pyarrow"
"""PyArrow implementation."""
PYSPARK = "pyspark"
"""PySpark implementation."""
POLARS = "polars"
"""Polars implementation."""
DASK = "dask"
"""Dask implementation."""
DUCKDB = "duckdb"
"""DuckDB implementation."""
IBIS = "ibis"
"""Ibis implementation."""
SQLFRAME = "sqlframe"
"""SQLFrame implementation."""
PYSPARK_CONNECT = "pyspark[connect]"
"""PySpark Connect implementation."""
UNKNOWN = "unknown"
"""Unknown implementation."""
def __str__(self) -> str:
return str(self.value)
@classmethod
def from_native_namespace(
cls: type[Self], native_namespace: ModuleType
) -> Implementation: # pragma: no cover
"""Instantiate Implementation object from a native namespace module.
Arguments:
native_namespace: Native namespace.
Returns:
Implementation.
"""
mapping = {
get_pandas(): Implementation.PANDAS,
get_modin(): Implementation.MODIN,
get_cudf(): Implementation.CUDF,
get_pyarrow(): Implementation.PYARROW,
get_pyspark_sql(): Implementation.PYSPARK,
get_polars(): Implementation.POLARS,
get_dask_dataframe(): Implementation.DASK,
get_duckdb(): Implementation.DUCKDB,
get_ibis(): Implementation.IBIS,
get_sqlframe(): Implementation.SQLFRAME,
get_pyspark_connect(): Implementation.PYSPARK_CONNECT,
}
return mapping.get(native_namespace, Implementation.UNKNOWN)
@classmethod
def from_string(
cls: type[Self], backend_name: str
) -> Implementation: # pragma: no cover
"""Instantiate Implementation object from a native namespace module.
Arguments:
backend_name: Name of backend, expressed as string.
Returns:
Implementation.
"""
try:
return cls(backend_name)
except ValueError:
return Implementation.UNKNOWN
@classmethod
def from_backend(
cls: type[Self], backend: str | Implementation | ModuleType
) -> Implementation:
"""Instantiate from native namespace module, string, or Implementation.
Arguments:
backend: Backend to instantiate Implementation from.
Returns:
Implementation.
"""
return (
cls.from_string(backend)
if isinstance(backend, str)
else backend
if isinstance(backend, Implementation)
else cls.from_native_namespace(backend)
)
def to_native_namespace(self) -> ModuleType:
"""Return the native namespace module corresponding to Implementation.
Returns:
Native module.
"""
if self is Implementation.UNKNOWN:
msg = "Cannot return native namespace from UNKNOWN Implementation"
raise AssertionError(msg)
self._backend_version()
module_name = _IMPLEMENTATION_TO_MODULE_NAME.get(self, self.value)
return _import_native_namespace(module_name)
def is_pandas(self) -> bool:
"""Return whether implementation is pandas.
Returns:
Boolean.
Examples:
>>> import pandas as pd
>>> import narwhals as nw
>>> df_native = pd.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_pandas()
True
"""
return self is Implementation.PANDAS
def is_pandas_like(self) -> bool:
"""Return whether implementation is pandas, Modin, or cuDF.
Returns:
Boolean.
Examples:
>>> import pandas as pd
>>> import narwhals as nw
>>> df_native = pd.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_pandas_like()
True
"""
return self in {Implementation.PANDAS, Implementation.MODIN, Implementation.CUDF}
def is_spark_like(self) -> bool:
"""Return whether implementation is pyspark or sqlframe.
Returns:
Boolean.
Examples:
>>> import pandas as pd
>>> import narwhals as nw
>>> df_native = pd.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_spark_like()
False
"""
return self in {
Implementation.PYSPARK,
Implementation.SQLFRAME,
Implementation.PYSPARK_CONNECT,
}
def is_polars(self) -> bool:
"""Return whether implementation is Polars.
Returns:
Boolean.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_polars()
True
"""
return self is Implementation.POLARS
def is_cudf(self) -> bool:
"""Return whether implementation is cuDF.
Returns:
Boolean.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_cudf()
False
"""
return self is Implementation.CUDF # pragma: no cover
def is_modin(self) -> bool:
"""Return whether implementation is Modin.
Returns:
Boolean.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_modin()
False
"""
return self is Implementation.MODIN # pragma: no cover
def is_pyspark(self) -> bool:
"""Return whether implementation is PySpark.
Returns:
Boolean.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_pyspark()
False
"""
return self is Implementation.PYSPARK # pragma: no cover
def is_pyspark_connect(self) -> bool:
"""Return whether implementation is PySpark.
Returns:
Boolean.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_pyspark_connect()
False
"""
return self is Implementation.PYSPARK_CONNECT # pragma: no cover
def is_pyarrow(self) -> bool:
"""Return whether implementation is PyArrow.
Returns:
Boolean.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_pyarrow()
False
"""
return self is Implementation.PYARROW # pragma: no cover
def is_dask(self) -> bool:
"""Return whether implementation is Dask.
Returns:
Boolean.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_dask()
False
"""
return self is Implementation.DASK # pragma: no cover
def is_duckdb(self) -> bool:
"""Return whether implementation is DuckDB.
Returns:
Boolean.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_duckdb()
False
"""
return self is Implementation.DUCKDB # pragma: no cover
def is_ibis(self) -> bool:
"""Return whether implementation is Ibis.
Returns:
Boolean.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_ibis()
False
"""
return self is Implementation.IBIS # pragma: no cover
def is_sqlframe(self) -> bool:
"""Return whether implementation is SQLFrame.
Returns:
Boolean.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3]})
>>> df = nw.from_native(df_native)
>>> df.implementation.is_sqlframe()
False
"""
return self is Implementation.SQLFRAME # pragma: no cover
def _backend_version(self) -> tuple[int, ...]:
"""Returns backend version."""
return backend_version(self)
MIN_VERSIONS: Mapping[Implementation, tuple[int, ...]] = {
Implementation.PANDAS: (1, 1, 3),
Implementation.MODIN: (0, 8, 2),
Implementation.CUDF: (24, 10),
Implementation.PYARROW: (13,),
Implementation.PYSPARK: (3, 5),
Implementation.PYSPARK_CONNECT: (3, 5),
Implementation.POLARS: (0, 20, 4),
Implementation.DASK: (2024, 8),
Implementation.DUCKDB: (1,),
Implementation.IBIS: (6,),
Implementation.SQLFRAME: (3, 22, 0),
}
_IMPLEMENTATION_TO_MODULE_NAME: Mapping[Implementation, str] = {
Implementation.DASK: "dask.dataframe",
Implementation.MODIN: "modin.pandas",
Implementation.PYSPARK: "pyspark.sql",
Implementation.PYSPARK_CONNECT: "pyspark.sql.connect",
}
"""Stores non default mapping from Implementation to module name"""
@lru_cache(maxsize=16)
def _import_native_namespace(module_name: str) -> ModuleType:
from importlib import import_module
return import_module(module_name)
# NOTE: We can safely use an unbounded cache, the size is constrained by `len(Implementation._member_names_)`
# Faster than `lru_cache`
# https://docs.python.org/3/library/functools.html#functools.cache
@cache
def backend_version(implementation: Implementation, /) -> tuple[int, ...]:
if not isinstance(implementation, Implementation):
assert_never(implementation)
if implementation is Implementation.UNKNOWN: # pragma: no cover
msg = "Cannot return backend version from UNKNOWN Implementation"
raise AssertionError(msg)
into_version: ModuleType | str
impl = implementation
module_name = _IMPLEMENTATION_TO_MODULE_NAME.get(impl, impl.value)
native_namespace = _import_native_namespace(module_name)
if impl.is_sqlframe():
import sqlframe._version
into_version = sqlframe._version
elif impl.is_pyspark() or impl.is_pyspark_connect(): # pragma: no cover
import pyspark # ignore-banned-import
into_version = pyspark
elif impl.is_dask():
import dask # ignore-banned-import
into_version = dask
else:
into_version = native_namespace
version = parse_version(into_version)
if version < (min_version := MIN_VERSIONS[impl]):
msg = f"Minimum version of {impl} supported by Narwhals is {min_version}, found: {version}"
raise ValueError(msg)
return version
def flatten(args: Any) -> list[Any]:
return list(args[0] if (len(args) == 1 and _is_iterable(args[0])) else args)
def tupleify(arg: Any) -> Any:
if not isinstance(arg, (list, tuple)): # pragma: no cover
return (arg,)
return arg
def _is_iterable(arg: Any | Iterable[Any]) -> bool:
from narwhals.series import Series
if (
(pd := get_pandas()) is not None and isinstance(arg, (pd.Series, pd.DataFrame))
) or (
(pl := get_polars()) is not None
and isinstance(arg, (pl.Series, pl.Expr, pl.DataFrame, pl.LazyFrame))
):
# Non-exhaustive check for common potential mistakes.
msg = (
f"Expected Narwhals class or scalar, got: {qualified_type_name(arg)!r}.\n\n"
"Hint: Perhaps you\n"
"- forgot a `nw.from_native` somewhere?\n"
"- used `pl.col` instead of `nw.col`?"
)
raise TypeError(msg)
return isinstance(arg, Iterable) and not isinstance(arg, (str, bytes, Series))
def parse_version(version: str | ModuleType | _SupportsVersion) -> tuple[int, ...]:
"""Simple version parser; split into a tuple of ints for comparison.
Arguments:
version: Version string, or object with one, to parse.
Returns:
Parsed version number.
"""
# lifted from Polars
# [marco]: Take care of DuckDB pre-releases which end with e.g. `-dev4108`
# and pandas pre-releases which end with e.g. .dev0+618.gb552dc95c9
version_str = version if isinstance(version, str) else version.__version__
version_str = re.sub(r"(\D?dev.*$)", "", version_str)
return tuple(int(re.sub(r"\D", "", v)) for v in version_str.split("."))
@overload
def isinstance_or_issubclass(
obj_or_cls: type, cls_or_tuple: type[_T]
) -> TypeIs[type[_T]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: object | type, cls_or_tuple: type[_T]
) -> TypeIs[_T | type[_T]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: type, cls_or_tuple: tuple[type[_T1], type[_T2]]
) -> TypeIs[type[_T1 | _T2]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: object | type, cls_or_tuple: tuple[type[_T1], type[_T2]]
) -> TypeIs[_T1 | _T2 | type[_T1 | _T2]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: type, cls_or_tuple: tuple[type[_T1], type[_T2], type[_T3]]
) -> TypeIs[type[_T1 | _T2 | _T3]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: object | type, cls_or_tuple: tuple[type[_T1], type[_T2], type[_T3]]
) -> TypeIs[_T1 | _T2 | _T3 | type[_T1 | _T2 | _T3]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: type, cls_or_tuple: tuple[type[_T1], type[_T2], type[_T3], type[_T4]]
) -> TypeIs[type[_T1 | _T2 | _T3 | _T4]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: object | type,
cls_or_tuple: tuple[type[_T1], type[_T2], type[_T3], type[_T4]],
) -> TypeIs[_T1 | _T2 | _T3 | _T4 | type[_T1 | _T2 | _T3 | _T4]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: type,
cls_or_tuple: tuple[type[_T1], type[_T2], type[_T3], type[_T4], type[_T5]],
) -> TypeIs[type[_T1 | _T2 | _T3 | _T4 | _T5]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: object | type,
cls_or_tuple: tuple[type[_T1], type[_T2], type[_T3], type[_T4], type[_T5]],
) -> TypeIs[_T1 | _T2 | _T3 | _T4 | _T5 | type[_T1 | _T2 | _T3 | _T4 | _T5]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: type,
cls_or_tuple: tuple[type[_T1], type[_T2], type[_T3], type[_T4], type[_T5], type[_T6]],
) -> TypeIs[type[_T1 | _T2 | _T3 | _T4 | _T5 | _T6]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: object | type,
cls_or_tuple: tuple[type[_T1], type[_T2], type[_T3], type[_T4], type[_T5], type[_T6]],
) -> TypeIs[
_T1 | _T2 | _T3 | _T4 | _T5 | _T6 | type[_T1 | _T2 | _T3 | _T4 | _T5 | _T6]
]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: type,
cls_or_tuple: tuple[
type[_T1], type[_T2], type[_T3], type[_T4], type[_T5], type[_T6], type[_T7]
],
) -> TypeIs[type[_T1 | _T2 | _T3 | _T4 | _T5 | _T6 | _T7]]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: object | type,
cls_or_tuple: tuple[
type[_T1], type[_T2], type[_T3], type[_T4], type[_T5], type[_T6], type[_T7]
],
) -> TypeIs[
_T1
| _T2
| _T3
| _T4
| _T5
| _T6
| _T7
| type[_T1 | _T2 | _T3 | _T4 | _T5 | _T6 | _T7]
]: ...
@overload
def isinstance_or_issubclass(
obj_or_cls: Any, cls_or_tuple: tuple[type, ...]
) -> TypeIs[Any]: ...
def isinstance_or_issubclass(obj_or_cls: Any, cls_or_tuple: Any) -> bool:
from narwhals.dtypes import DType
if isinstance(obj_or_cls, DType):
return isinstance(obj_or_cls, cls_or_tuple)
return isinstance(obj_or_cls, cls_or_tuple) or (
isinstance(obj_or_cls, type) and issubclass(obj_or_cls, cls_or_tuple)
)
def validate_laziness(items: Iterable[Any]) -> None:
from narwhals.dataframe import DataFrame, LazyFrame
if all(isinstance(item, DataFrame) for item in items) or (
all(isinstance(item, LazyFrame) for item in items)
):
return
msg = f"The items to concatenate should either all be eager, or all lazy, got: {[type(item) for item in items]}"
raise TypeError(msg)
def maybe_align_index(
lhs: FrameOrSeriesT, rhs: Series[Any] | DataFrame[Any] | LazyFrame[Any]
) -> FrameOrSeriesT:
"""Align `lhs` to the Index of `rhs`, if they're both pandas-like.
Arguments:
lhs: Dataframe or Series.
rhs: Dataframe or Series to align with.
Returns:
Same type as input.
Notes:
This is only really intended for backwards-compatibility purposes,
for example if your library already aligns indices for users.
If you're designing a new library, we highly encourage you to not
rely on the Index.
For non-pandas-like inputs, this only checks that `lhs` and `rhs`
are the same length.
Examples:
>>> import pandas as pd
>>> import polars as pl
>>> import narwhals as nw
>>> df_pd = pd.DataFrame({"a": [1, 2]}, index=[3, 4])
>>> s_pd = pd.Series([6, 7], index=[4, 3])
>>> df = nw.from_native(df_pd)
>>> s = nw.from_native(s_pd, series_only=True)
>>> nw.to_native(nw.maybe_align_index(df, s))
a
4 2
3 1
"""
from narwhals._pandas_like.dataframe import PandasLikeDataFrame
from narwhals._pandas_like.series import PandasLikeSeries
def _validate_index(index: Any) -> None:
if not index.is_unique:
msg = "given index doesn't have a unique index"
raise ValueError(msg)
lhs_any = cast("Any", lhs)
rhs_any = cast("Any", rhs)
if isinstance(
getattr(lhs_any, "_compliant_frame", None), PandasLikeDataFrame
) and isinstance(getattr(rhs_any, "_compliant_frame", None), PandasLikeDataFrame):
_validate_index(lhs_any._compliant_frame.native.index)
_validate_index(rhs_any._compliant_frame.native.index)
return lhs_any._with_compliant(
lhs_any._compliant_frame._with_native(
lhs_any._compliant_frame.native.loc[rhs_any._compliant_frame.native.index]
)
)
if isinstance(
getattr(lhs_any, "_compliant_frame", None), PandasLikeDataFrame
) and isinstance(getattr(rhs_any, "_compliant_series", None), PandasLikeSeries):
_validate_index(lhs_any._compliant_frame.native.index)
_validate_index(rhs_any._compliant_series.native.index)
return lhs_any._with_compliant(
lhs_any._compliant_frame._with_native(
lhs_any._compliant_frame.native.loc[
rhs_any._compliant_series.native.index
]
)
)
if isinstance(
getattr(lhs_any, "_compliant_series", None), PandasLikeSeries
) and isinstance(getattr(rhs_any, "_compliant_frame", None), PandasLikeDataFrame):
_validate_index(lhs_any._compliant_series.native.index)
_validate_index(rhs_any._compliant_frame.native.index)
return lhs_any._with_compliant(
lhs_any._compliant_series._with_native(
lhs_any._compliant_series.native.loc[
rhs_any._compliant_frame.native.index
]
)
)
if isinstance(
getattr(lhs_any, "_compliant_series", None), PandasLikeSeries
) and isinstance(getattr(rhs_any, "_compliant_series", None), PandasLikeSeries):
_validate_index(lhs_any._compliant_series.native.index)
_validate_index(rhs_any._compliant_series.native.index)
return lhs_any._with_compliant(
lhs_any._compliant_series._with_native(
lhs_any._compliant_series.native.loc[
rhs_any._compliant_series.native.index
]
)
)
if len(lhs_any) != len(rhs_any):
msg = f"Expected `lhs` and `rhs` to have the same length, got {len(lhs_any)} and {len(rhs_any)}"
raise ValueError(msg)
return lhs
def maybe_get_index(obj: DataFrame[Any] | LazyFrame[Any] | Series[Any]) -> Any | None:
"""Get the index of a DataFrame or a Series, if it's pandas-like.
Arguments:
obj: Dataframe or Series.
Returns:
Same type as input.
Notes:
This is only really intended for backwards-compatibility purposes,
for example if your library already aligns indices for users.
If you're designing a new library, we highly encourage you to not
rely on the Index.
For non-pandas-like inputs, this returns `None`.
Examples:
>>> import pandas as pd
>>> import polars as pl
>>> import narwhals as nw
>>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]})
>>> df = nw.from_native(df_pd)
>>> nw.maybe_get_index(df)
RangeIndex(start=0, stop=2, step=1)
>>> series_pd = pd.Series([1, 2])
>>> series = nw.from_native(series_pd, series_only=True)
>>> nw.maybe_get_index(series)
RangeIndex(start=0, stop=2, step=1)
"""
obj_any = cast("Any", obj)
native_obj = obj_any.to_native()
if is_pandas_like_dataframe(native_obj) or is_pandas_like_series(native_obj):
return native_obj.index
return None
def maybe_set_index(
obj: FrameOrSeriesT,
column_names: str | list[str] | None = None,
*,
index: Series[IntoSeriesT] | list[Series[IntoSeriesT]] | None = None,
) -> FrameOrSeriesT:
"""Set the index of a DataFrame or a Series, if it's pandas-like.
Arguments:
obj: object for which maybe set the index (can be either a Narwhals `DataFrame`
or `Series`).
column_names: name or list of names of the columns to set as index.
For dataframes, only one of `column_names` and `index` can be specified but
not both. If `column_names` is passed and `df` is a Series, then a
`ValueError` is raised.
index: series or list of series to set as index.
Returns:
Same type as input.
Raises:
ValueError: If one of the following conditions happens
- none of `column_names` and `index` are provided
- both `column_names` and `index` are provided
- `column_names` is provided and `df` is a Series
Notes:
This is only really intended for backwards-compatibility purposes, for example if
your library already aligns indices for users.
If you're designing a new library, we highly encourage you to not
rely on the Index.
For non-pandas-like inputs, this is a no-op.
Examples:
>>> import pandas as pd
>>> import polars as pl
>>> import narwhals as nw
>>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]})
>>> df = nw.from_native(df_pd)
>>> nw.to_native(nw.maybe_set_index(df, "b")) # doctest: +NORMALIZE_WHITESPACE
a
b
4 1
5 2
"""
from narwhals.translate import to_native
df_any = cast("Any", obj)
native_obj = df_any.to_native()
if column_names is not None and index is not None:
msg = "Only one of `column_names` or `index` should be provided"
raise ValueError(msg)
if not column_names and index is None:
msg = "Either `column_names` or `index` should be provided"
raise ValueError(msg)
if index is not None:
keys = (
[to_native(idx, pass_through=True) for idx in index]
if _is_iterable(index)
else to_native(index, pass_through=True)
)
else:
keys = column_names
if is_pandas_like_dataframe(native_obj):
return df_any._with_compliant(
df_any._compliant_frame._with_native(native_obj.set_index(keys))
)
elif is_pandas_like_series(native_obj):
from narwhals._pandas_like.utils import set_index
if column_names:
msg = "Cannot set index using column names on a Series"
raise ValueError(msg)
native_obj = set_index(
native_obj,
keys,
implementation=obj._compliant_series._implementation, # type: ignore[union-attr]
)
return df_any._with_compliant(df_any._compliant_series._with_native(native_obj))
else:
return df_any
def maybe_reset_index(obj: FrameOrSeriesT) -> FrameOrSeriesT:
"""Reset the index to the default integer index of a DataFrame or a Series, if it's pandas-like.
Arguments:
obj: Dataframe or Series.
Returns:
Same type as input.
Notes:
This is only really intended for backwards-compatibility purposes,
for example if your library already resets the index for users.
If you're designing a new library, we highly encourage you to not
rely on the Index.
For non-pandas-like inputs, this is a no-op.
Examples:
>>> import pandas as pd
>>> import polars as pl
>>> import narwhals as nw
>>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]}, index=([6, 7]))
>>> df = nw.from_native(df_pd)
>>> nw.to_native(nw.maybe_reset_index(df))
a b
0 1 4
1 2 5
>>> series_pd = pd.Series([1, 2])
>>> series = nw.from_native(series_pd, series_only=True)
>>> nw.maybe_get_index(series)
RangeIndex(start=0, stop=2, step=1)
"""
obj_any = cast("Any", obj)
native_obj = obj_any.to_native()
if is_pandas_like_dataframe(native_obj):
native_namespace = obj_any.__native_namespace__()
if _has_default_index(native_obj, native_namespace):
return obj_any
return obj_any._with_compliant(
obj_any._compliant_frame._with_native(native_obj.reset_index(drop=True))
)
if is_pandas_like_series(native_obj):
native_namespace = obj_any.__native_namespace__()
if _has_default_index(native_obj, native_namespace):
return obj_any
return obj_any._with_compliant(
obj_any._compliant_series._with_native(native_obj.reset_index(drop=True))
)
return obj_any
def _is_range_index(obj: Any, native_namespace: Any) -> TypeIs[pd.RangeIndex]:
return isinstance(obj, native_namespace.RangeIndex)
def _has_default_index(
native_frame_or_series: pd.Series[Any] | pd.DataFrame, native_namespace: Any
) -> bool:
index = native_frame_or_series.index
return (
_is_range_index(index, native_namespace)
and index.start == 0
and index.stop == len(index)
and index.step == 1
)
def maybe_convert_dtypes(
obj: FrameOrSeriesT, *args: bool, **kwargs: bool | str
) -> FrameOrSeriesT:
"""Convert columns or series to the best possible dtypes using dtypes supporting ``pd.NA``, if df is pandas-like.
Arguments:
obj: DataFrame or Series.
*args: Additional arguments which gets passed through.
**kwargs: Additional arguments which gets passed through.
Returns:
Same type as input.
Notes:
For non-pandas-like inputs, this is a no-op.
Also, `args` and `kwargs` just get passed down to the underlying library as-is.
Examples:
>>> import pandas as pd
>>> import polars as pl
>>> import narwhals as nw
>>> import numpy as np
>>> df_pd = pd.DataFrame(
... {
... "a": pd.Series([1, 2, 3], dtype=np.dtype("int32")),
... "b": pd.Series([True, False, np.nan], dtype=np.dtype("O")),
... }
... )
>>> df = nw.from_native(df_pd)
>>> nw.to_native(
... nw.maybe_convert_dtypes(df)
... ).dtypes # doctest: +NORMALIZE_WHITESPACE
a Int32
b boolean
dtype: object
"""
obj_any = cast("Any", obj)
native_obj = obj_any.to_native()
if is_pandas_like_dataframe(native_obj):
return obj_any._with_compliant(
obj_any._compliant_frame._with_native(
native_obj.convert_dtypes(*args, **kwargs)
)
)
if is_pandas_like_series(native_obj):
return obj_any._with_compliant(
obj_any._compliant_series._with_native(
native_obj.convert_dtypes(*args, **kwargs)
)
)
return obj_any
def scale_bytes(sz: int, unit: SizeUnit) -> int | float:
"""Scale size in bytes to other size units (eg: "kb", "mb", "gb", "tb").
Arguments:
sz: original size in bytes
unit: size unit to convert into
Returns:
Integer or float.
"""
if unit in {"b", "bytes"}:
return sz
elif unit in {"kb", "kilobytes"}:
return sz / 1024
elif unit in {"mb", "megabytes"}:
return sz / 1024**2
elif unit in {"gb", "gigabytes"}:
return sz / 1024**3
elif unit in {"tb", "terabytes"}:
return sz / 1024**4
else:
msg = f"`unit` must be one of {{'b', 'kb', 'mb', 'gb', 'tb'}}, got {unit!r}"
raise ValueError(msg)
def is_ordered_categorical(series: Series[Any]) -> bool:
"""Return whether indices of categories are semantically meaningful.
This is a convenience function to accessing what would otherwise be
the `is_ordered` property from the DataFrame Interchange Protocol,
see https://data-apis.org/dataframe-protocol/latest/API.html.
- For Polars:
- Enums are always ordered.
- Categoricals are ordered if `dtype.ordering == "physical"`.
- For pandas-like APIs:
- Categoricals are ordered if `dtype.cat.ordered == True`.
- For PyArrow table:
- Categoricals are ordered if `dtype.type.ordered == True`.
Arguments:
series: Input Series.
Returns:
Whether the Series is an ordered categorical.
Examples:
>>> import narwhals as nw
>>> import pandas as pd
>>> import polars as pl
>>> data = ["x", "y"]
>>> s_pd = pd.Series(data, dtype=pd.CategoricalDtype(ordered=True))
>>> s_pl = pl.Series(data, dtype=pl.Categorical(ordering="physical"))
Let's define a library-agnostic function:
>>> @nw.narwhalify
... def func(s):
... return nw.is_ordered_categorical(s)
Then, we can pass any supported library to `func`:
>>> func(s_pd)
True
>>> func(s_pl)
True
"""
from narwhals._interchange.series import InterchangeSeries
dtypes = series._compliant_series._version.dtypes
compliant = series._compliant_series
# If it doesn't match any branches, let's just play it safe and return False.
result: bool = False
if isinstance(compliant, InterchangeSeries) and isinstance(
series.dtype, dtypes.Categorical
):
result = compliant.native.describe_categorical["is_ordered"]
elif series.dtype == dtypes.Enum:
result = True
elif series.dtype != dtypes.Categorical:
result = False
else:
native = series.to_native()
if is_polars_series(native):
result = cast("pl.Categorical", native.dtype).ordering == "physical"
elif is_pandas_like_series(native):
result = bool(native.cat.ordered)
elif is_pyarrow_chunked_array(native):
from narwhals._arrow.utils import is_dictionary
result = is_dictionary(native.type) and native.type.ordered
return result
def generate_unique_token(
n_bytes: int, columns: Container[str]
) -> str: # pragma: no cover
msg = (
"Use `generate_temporary_column_name` instead. `generate_unique_token` is "
"deprecated and it will be removed in future versions"
)
issue_deprecation_warning(msg, _version="1.13.0")
return generate_temporary_column_name(n_bytes=n_bytes, columns=columns)
def generate_temporary_column_name(n_bytes: int, columns: Container[str]) -> str:
"""Generates a unique column name that is not present in the given list of columns.
It relies on [python secrets token_hex](https://docs.python.org/3/library/secrets.html#secrets.token_hex)
function to return a string nbytes random bytes.
Arguments:
n_bytes: The number of bytes to generate for the token.
columns: The list of columns to check for uniqueness.
Returns:
A unique token that is not present in the given list of columns.
Raises:
AssertionError: If a unique token cannot be generated after 100 attempts.
Examples:
>>> import narwhals as nw
>>> columns = ["abc", "xyz"]
>>> nw.generate_temporary_column_name(n_bytes=8, columns=columns) not in columns
True
"""
counter = 0
while True:
# Prepend `'nw'` to ensure it always starts with a character
# https://github.com/narwhals-dev/narwhals/issues/2510
token = f"nw{token_hex(n_bytes - 1)}"
if token not in columns:
return token
counter += 1
if counter > 100:
msg = (
"Internal Error: Narwhals was not able to generate a column name with "
f"{n_bytes=} and not in {columns}"
)
raise AssertionError(msg)
def parse_columns_to_drop(
frame: _StoresColumns, subset: Iterable[str], /, *, strict: bool
) -> list[str]:
if not strict:
return list(set(frame.columns).intersection(subset))
to_drop = list(subset)
if error := check_columns_exist(to_drop, available=frame.columns):
raise error
return to_drop
def is_sequence_but_not_str(sequence: Sequence[_T] | Any) -> TypeIs[Sequence[_T]]:
return isinstance(sequence, Sequence) and not isinstance(sequence, str)
def is_slice_none(obj: Any) -> TypeIs[_SliceNone]:
return isinstance(obj, slice) and obj == slice(None)
def is_sized_multi_index_selector(
obj: Any,
) -> TypeIs[SizedMultiIndexSelector[Series[Any] | CompliantSeries[Any]]]:
return (
(
is_sequence_but_not_str(obj)
and ((len(obj) > 0 and isinstance(obj[0], int)) or (len(obj) == 0))
)
or is_numpy_array_1d_int(obj)
or is_narwhals_series_int(obj)
or is_compliant_series_int(obj)
)
def is_sequence_like(
obj: Sequence[_T] | Any,
) -> TypeIs[Sequence[_T] | Series[Any] | _1DArray]:
return (
is_sequence_but_not_str(obj)
or is_numpy_array_1d(obj)
or is_narwhals_series(obj)
or is_compliant_series(obj)
)
def is_slice_index(obj: Any) -> TypeIs[_SliceIndex]:
return isinstance(obj, slice) and (
isinstance(obj.start, int)
or isinstance(obj.stop, int)
or (isinstance(obj.step, int) and obj.start is None and obj.stop is None)
)
def is_range(obj: Any) -> TypeIs[range]:
return isinstance(obj, range)
def is_single_index_selector(obj: Any) -> TypeIs[SingleIndexSelector]:
return bool(isinstance(obj, int) and not isinstance(obj, bool))
def is_index_selector(
obj: Any,
) -> TypeIs[SingleIndexSelector | MultiIndexSelector[Series[Any] | CompliantSeries[Any]]]:
return (
is_single_index_selector(obj)
or is_sized_multi_index_selector(obj)
or is_slice_index(obj)
)
def is_list_of(obj: Any, tp: type[_T]) -> TypeIs[list[_T]]:
# Check if an object is a list of `tp`, only sniffing the first element.
return bool(isinstance(obj, list) and obj and isinstance(obj[0], tp))
def is_sequence_of(obj: Any, tp: type[_T]) -> TypeIs[Sequence[_T]]:
# Check if an object is a sequence of `tp`, only sniffing the first element.
return bool(
is_sequence_but_not_str(obj)
and (first := next(iter(obj), None))
and isinstance(first, tp)
)
def validate_strict_and_pass_though(
strict: bool | None, # noqa: FBT001
pass_through: bool | None, # noqa: FBT001
*,
pass_through_default: bool,
) -> bool:
if strict is None and pass_through is None:
pass_through = pass_through_default
elif strict is not None and pass_through is None:
pass_through = not strict
elif strict is None and pass_through is not None:
pass
else:
msg = "Cannot pass both `strict` and `pass_through`"
raise ValueError(msg)
return pass_through
def deprecate_native_namespace(
*, warn_version: str = "", required: bool = False
) -> Callable[[Callable[P, R]], Callable[P, R]]:
"""Decorator to transition from `native_namespace` to `backend` argument.
Arguments:
warn_version: Emit a deprecation warning from this version.
required: Raise when both `native_namespace`, `backend` are `None`.
Returns:
Wrapped function, with `native_namespace` **removed**.
"""
def decorate(fn: Callable[P, R], /) -> Callable[P, R]:
@wraps(fn)
def wrapper(*args: P.args, **kwds: P.kwargs) -> R:
backend = kwds.pop("backend", None)
native_namespace = kwds.pop("native_namespace", None)
if native_namespace is not None and backend is None:
if warn_version:
msg = (
"`native_namespace` is deprecated, please use `backend` instead.\n\n"
"Note: `native_namespace` will remain available in `narwhals.stable.v1`.\n"
"See https://narwhals-dev.github.io/narwhals/backcompat/ for more information.\n"
)
issue_deprecation_warning(msg, _version=warn_version)
backend = native_namespace
elif native_namespace is not None and backend is not None:
msg = "Can't pass both `native_namespace` and `backend`"
raise ValueError(msg)
elif native_namespace is None and backend is None and required:
msg = f"`backend` must be specified in `{fn.__name__}`."
raise ValueError(msg)
kwds["backend"] = backend
return fn(*args, **kwds)
return wrapper
return decorate
def _validate_rolling_arguments(
window_size: int, min_samples: int | None
) -> tuple[int, int]:
ensure_type(window_size, int, param_name="window_size")
ensure_type(min_samples, int, type(None), param_name="min_samples")
if window_size < 1:
msg = "window_size must be greater or equal than 1"
raise ValueError(msg)
if min_samples is not None:
if min_samples < 1:
msg = "min_samples must be greater or equal than 1"
raise ValueError(msg)
if min_samples > window_size:
msg = "`min_samples` must be less or equal than `window_size`"
raise InvalidOperationError(msg)
else:
min_samples = window_size
return window_size, min_samples
def generate_repr(header: str, native_repr: str) -> str:
try:
terminal_width = os.get_terminal_size().columns
except OSError:
terminal_width = int(os.getenv("COLUMNS", 80)) # noqa: PLW1508
native_lines = native_repr.expandtabs().splitlines()
max_native_width = max(len(line) for line in native_lines)
if max_native_width + 2 <= terminal_width:
length = max(max_native_width, len(header))
output = f"┌{'─' * length}┐\n"
header_extra = length - len(header)
output += f"|{' ' * (header_extra // 2)}{header}{' ' * (header_extra // 2 + header_extra % 2)}|\n"
output += f"|{'-' * (length)}|\n"
start_extra = (length - max_native_width) // 2
end_extra = (length - max_native_width) // 2 + (length - max_native_width) % 2
for line in native_lines:
output += f"|{' ' * (start_extra)}{line}{' ' * (end_extra + max_native_width - len(line))}|\n"
output += f"└{'─' * length}┘"
return output
diff = 39 - len(header)
return (
f"┌{'─' * (39)}┐\n"
f"|{' ' * (diff // 2)}{header}{' ' * (diff // 2 + diff % 2)}|\n"
"| Use `.to_native` to see native output |\n└"
f"{'─' * 39}┘"
)
def check_columns_exist(
subset: Collection[str], /, *, available: Collection[str]
) -> ColumnNotFoundError | None:
if missing := set(subset).difference(available):
return ColumnNotFoundError.from_missing_and_available_column_names(
missing, available
)
return None
def check_column_names_are_unique(columns: Collection[str]) -> None:
if len(columns) != len(set(columns)):
from collections import Counter
counter = Counter(columns)
duplicates = {k: v for k, v in counter.items() if v > 1}
msg = "".join(f"\n- '{k}' {v} times" for k, v in duplicates.items())
msg = f"Expected unique column names, got:{msg}"
raise DuplicateError(msg)
def _parse_time_unit_and_time_zone(
time_unit: TimeUnit | Iterable[TimeUnit] | None,
time_zone: str | timezone | Iterable[str | timezone | None] | None,
) -> tuple[Set[TimeUnit], Set[str | None]]:
time_units: Set[TimeUnit] = (
{"ms", "us", "ns", "s"}
if time_unit is None
else {time_unit}
if isinstance(time_unit, str)
else set(time_unit)
)
time_zones: Set[str | None] = (
{None}
if time_zone is None
else {str(time_zone)}
if isinstance(time_zone, (str, timezone))
else {str(tz) if tz is not None else None for tz in time_zone}
)
return time_units, time_zones
def dtype_matches_time_unit_and_time_zone(
dtype: DType, dtypes: DTypes, time_units: Set[TimeUnit], time_zones: Set[str | None]
) -> bool:
return (
isinstance(dtype, dtypes.Datetime)
and (dtype.time_unit in time_units)
and (
dtype.time_zone in time_zones
or ("*" in time_zones and dtype.time_zone is not None)
)
)
def get_column_names(frame: _StoresColumns, /) -> Sequence[str]:
return frame.columns
def exclude_column_names(frame: _StoresColumns, names: Container[str]) -> Sequence[str]:
return [col_name for col_name in frame.columns if col_name not in names]
def passthrough_column_names(names: Sequence[str], /) -> EvalNames[Any]:
def fn(_frame: Any, /) -> Sequence[str]:
return names
return fn
def _hasattr_static(obj: Any, attr: str) -> bool:
sentinel = object()
return getattr_static(obj, attr, sentinel) is not sentinel
def is_compliant_dataframe(
obj: CompliantDataFrame[
CompliantSeriesT, CompliantExprT, NativeFrameT_co, ToNarwhalsT_co
]
| Any,
) -> TypeIs[
CompliantDataFrame[CompliantSeriesT, CompliantExprT, NativeFrameT_co, ToNarwhalsT_co]
]:
return _hasattr_static(obj, "__narwhals_dataframe__")
def is_compliant_lazyframe(
obj: CompliantLazyFrame[CompliantExprT, NativeFrameT_co, ToNarwhalsT_co] | Any,
) -> TypeIs[CompliantLazyFrame[CompliantExprT, NativeFrameT_co, ToNarwhalsT_co]]:
return _hasattr_static(obj, "__narwhals_lazyframe__")
def is_compliant_series(
obj: CompliantSeries[NativeSeriesT_co] | Any,
) -> TypeIs[CompliantSeries[NativeSeriesT_co]]:
return _hasattr_static(obj, "__narwhals_series__")
def is_compliant_series_int(
obj: CompliantSeries[NativeSeriesT_co] | Any,
) -> TypeIs[CompliantSeries[NativeSeriesT_co]]:
return is_compliant_series(obj) and obj.dtype.is_integer()
def is_compliant_expr(
obj: CompliantExpr[CompliantFrameT, CompliantSeriesOrNativeExprT_co] | Any,
) -> TypeIs[CompliantExpr[CompliantFrameT, CompliantSeriesOrNativeExprT_co]]:
return hasattr(obj, "__narwhals_expr__")
def is_eager_allowed(obj: Implementation) -> TypeIs[EagerAllowedImplementation]:
return obj in {
Implementation.PANDAS,
Implementation.MODIN,
Implementation.CUDF,
Implementation.POLARS,
Implementation.PYARROW,
}
def has_native_namespace(obj: Any) -> TypeIs[SupportsNativeNamespace]:
return _hasattr_static(obj, "__native_namespace__")
def _supports_dataframe_interchange(obj: Any) -> TypeIs[DataFrameLike]:
return hasattr(obj, "__dataframe__")
def supports_arrow_c_stream(obj: Any) -> TypeIs[ArrowStreamExportable]:
return _hasattr_static(obj, "__arrow_c_stream__")
def _remap_full_join_keys(
left_on: Collection[str], right_on: Collection[str], suffix: str
) -> dict[str, str]:
"""Remap join keys to avoid collisions.
If left keys collide with the right keys, append the suffix.
If there's no collision, let the right keys be.
Arguments:
left_on: Left keys.
right_on: Right keys.
suffix: Suffix to append to right keys.
Returns:
A map of old to new right keys.
"""
right_keys_suffixed = (
f"{key}{suffix}" if key in left_on else key for key in right_on
)
return dict(zip(right_on, right_keys_suffixed))
def _into_arrow_table(data: IntoArrowTable, context: _LimitedContext, /) -> pa.Table:
"""Guards `ArrowDataFrame.from_arrow` w/ safer imports.
Arguments:
data: Object which implements `__arrow_c_stream__`.
context: Initialized compliant object.
Returns:
A PyArrow Table.
"""
if find_spec("pyarrow"):
ns = context._version.namespace.from_backend("pyarrow").compliant
return ns._dataframe.from_arrow(data, context=ns).native
else: # pragma: no cover
msg = f"'pyarrow>=14.0.0' is required for `from_arrow` for object of type {qualified_type_name(data)!r}."
raise ModuleNotFoundError(msg)
# TODO @dangotbanned: Extend with runtime behavior for `v1.*`
# See `narwhals.exceptions.NarwhalsUnstableWarning`
def unstable(fn: _Fn, /) -> _Fn:
"""Visual-only marker for unstable functionality.
Arguments:
fn: Function to decorate.
Returns:
Decorated function (unchanged).
Examples:
>>> from narwhals._utils import unstable
>>> @unstable
... def a_work_in_progress_feature(*args):
... return args
>>>
>>> a_work_in_progress_feature.__name__
'a_work_in_progress_feature'
>>> a_work_in_progress_feature(1, 2, 3)
(1, 2, 3)
"""
return fn
def _is_naive_format(format: str) -> bool:
"""Determines if a datetime format string is 'naive', i.e., does not include timezone information.
A format is considered naive if it does not contain any of the following
- '%s': Unix timestamp
- '%z': UTC offset
- 'Z' : UTC timezone designator
Arguments:
format: The datetime format string to check.
Returns:
bool: True if the format is naive (does not include timezone info), False otherwise.
"""
return not any(x in format for x in ("%s", "%z", "Z"))
class not_implemented: # noqa: N801
"""Mark some functionality as unsupported.
Arguments:
alias: optional name used instead of the data model hook [`__set_name__`].
Returns:
An exception-raising [descriptor].
Notes:
- Attribute/method name *doesn't* need to be declared twice
- Allows different behavior when looked up on the class vs instance
- Allows us to use `isinstance(...)` instead of monkeypatching an attribute to the function
Examples:
>>> from narwhals._utils import not_implemented
>>> class Thing:
... def totally_ready(self) -> str:
... return "I'm ready!"
...
... not_ready_yet = not_implemented()
>>>
>>> thing = Thing()
>>> thing.totally_ready()
"I'm ready!"
>>> thing.not_ready_yet()
Traceback (most recent call last):
...
NotImplementedError: 'not_ready_yet' is not implemented for: 'Thing'.
...
>>> isinstance(Thing.not_ready_yet, not_implemented)
True
[`__set_name__`]: https://docs.python.org/3/reference/datamodel.html#object.__set_name__
[descriptor]: https://docs.python.org/3/howto/descriptor.html
"""
def __init__(self, alias: str | None = None, /) -> None:
# NOTE: Don't like this
# Trying to workaround `mypy` requiring `@property` everywhere
self._alias: str | None = alias
def __repr__(self) -> str:
return f"<{type(self).__name__}>: {self._name_owner}.{self._name}"
def __set_name__(self, owner: type[_T], name: str) -> None:
# https://docs.python.org/3/howto/descriptor.html#customized-names
self._name_owner: str = owner.__name__
self._name: str = self._alias or name
def __get__(
self, instance: _T | Literal["raise"] | None, owner: type[_T] | None = None, /
) -> Any:
if instance is None:
# NOTE: Branch for `cls._name`
# We can check that to see if an instance of `type(self)` for
# https://narwhals-dev.github.io/narwhals/api-completeness/expr/
return self
# NOTE: Prefer not exposing the actual class we're defining in
# `_implementation` may not be available everywhere
who = getattr(instance, "_implementation", self._name_owner)
_raise_not_implemented_error(self._name, who)
return None # pragma: no cover
def __call__(self, *args: Any, **kwds: Any) -> Any:
# NOTE: Purely to duck-type as assignable to **any** instance method
# Wouldn't be reachable through *regular* attribute access
return self.__get__("raise")
@classmethod
def deprecated(cls, message: LiteralString, /) -> Self:
"""Alt constructor, wraps with `@deprecated`.
Arguments:
message: **Static-only** deprecation message, emitted in an IDE.
Returns:
An exception-raising [descriptor].
[descriptor]: https://docs.python.org/3/howto/descriptor.html
"""
obj = cls()
return deprecated(message)(obj)
def _raise_not_implemented_error(what: str, who: str, /) -> NotImplementedError:
msg = (
f"{what!r} is not implemented for: {who!r}.\n\n"
"If you would like to see this functionality in `narwhals`, "
"please open an issue at: https://github.com/narwhals-dev/narwhals/issues"
)
raise NotImplementedError(msg)
class requires: # noqa: N801
"""Method decorator for raising under certain constraints.
Attributes:
_min_version: Minimum backend version.
_hint: Optional suggested alternative.
Examples:
>>> from narwhals._utils import requires, Implementation
>>> class SomeBackend:
... _implementation = Implementation.PYARROW
... _backend_version = 20, 0, 0
...
... @requires.backend_version((9000, 0, 0))
... def really_complex_feature(self) -> str:
... return "hello"
>>> backend = SomeBackend()
>>> backend.really_complex_feature()
Traceback (most recent call last):
...
NotImplementedError: `really_complex_feature` is only available in 'pyarrow>=9000.0.0', found version '20.0.0'.
"""
_min_version: tuple[int, ...]
_hint: str
@classmethod
def backend_version(cls, minimum: tuple[int, ...], /, hint: str = "") -> Self:
"""Method decorator for raising below a minimum `_backend_version`.
Arguments:
minimum: Minimum backend version.
hint: Optional suggested alternative.
Returns:
An exception-raising decorator.
"""
obj = cls.__new__(cls)
obj._min_version = minimum
obj._hint = hint
return obj
@staticmethod
def _unparse_version(backend_version: tuple[int, ...], /) -> str:
return ".".join(f"{d}" for d in backend_version)
def _ensure_version(self, instance: _FullContext, /) -> None:
if instance._backend_version >= self._min_version:
return
method = self._wrapped_name
backend = instance._implementation
minimum = self._unparse_version(self._min_version)
found = self._unparse_version(instance._backend_version)
msg = f"`{method}` is only available in '{backend}>={minimum}', found version {found!r}."
if self._hint:
msg = f"{msg}\n{self._hint}"
raise NotImplementedError(msg)
def __call__(self, fn: _Method[_ContextT, P, R], /) -> _Method[_ContextT, P, R]:
self._wrapped_name = fn.__name__
@wraps(fn)
def wrapper(instance: _ContextT, *args: P.args, **kwds: P.kwargs) -> R:
self._ensure_version(instance)
return fn(instance, *args, **kwds)
# NOTE: Only getting a complaint from `mypy`
return wrapper # type: ignore[return-value]
def convert_str_slice_to_int_slice(
str_slice: _SliceName, columns: Sequence[str]
) -> tuple[int | None, int | None, Any]:
start = columns.index(str_slice.start) if str_slice.start is not None else None
stop = columns.index(str_slice.stop) + 1 if str_slice.stop is not None else None
step = str_slice.step
return (start, stop, step)
def inherit_doc(
tp_parent: Callable[P, R1], /
) -> Callable[[_Constructor[_T, P, R2]], _Constructor[_T, P, R2]]:
"""Steal the class-level docstring from parent and attach to child `__init__`.
Returns:
Decorated constructor.
Notes:
- Passes static typing (mostly)
- Passes at runtime
"""
def decorate(init_child: _Constructor[_T, P, R2], /) -> _Constructor[_T, P, R2]:
if init_child.__name__ == "__init__" and issubclass(type(tp_parent), type):
init_child.__doc__ = getdoc(tp_parent)
return init_child
else: # pragma: no cover
msg = (
f"`@{inherit_doc.__name__}` is only allowed to decorate an `__init__` with a class-level doc.\n"
f"Method: {init_child.__qualname__!r}\n"
f"Parent: {tp_parent!r}"
)
raise TypeError(msg)
return decorate
def qualified_type_name(obj: object | type[Any], /) -> str:
tp = obj if isinstance(obj, type) else type(obj)
module = tp.__module__ if tp.__module__ != "builtins" else ""
return f"{module}.{tp.__name__}".lstrip(".")
def ensure_type(obj: Any, /, *valid_types: type[Any], param_name: str = "") -> None:
"""Validate that an object is an instance of one or more specified types.
Parameters:
obj: The object to validate.
*valid_types: One or more valid types that `obj` is expected to match.
param_name: The name of the parameter being validated.
Used to improve error message clarity.
Raises:
TypeError: If `obj` is not an instance of any of the provided `valid_types`.
Examples:
>>> from narwhals._utils import ensure_type
>>> ensure_type(42, int, float)
>>> ensure_type("hello", str)
>>> ensure_type("hello", int, param_name="test")
Traceback (most recent call last):
...
TypeError: Expected 'int', got: 'str'
test='hello'
^^^^^^^
>>> import polars as pl
>>> import pandas as pd
>>> df = pl.DataFrame([[1], [2], [3], [4], [5]], schema=[*"abcde"])
>>> ensure_type(df, pd.DataFrame, param_name="df")
Traceback (most recent call last):
...
TypeError: Expected 'pandas.core.frame.DataFrame', got: 'polars.dataframe.frame.DataFrame'
df=polars.dataframe.frame.DataFrame(...)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"""
if not isinstance(obj, valid_types): # pragma: no cover
tp_names = " | ".join(qualified_type_name(tp) for tp in valid_types)
msg = f"Expected {tp_names!r}, got: {qualified_type_name(obj)!r}"
if param_name:
left_pad = " " * 4
val = repr(obj)
if len(val) > 40: # truncate long reprs
val = f"{qualified_type_name(obj)}(...)"
assign = f"{left_pad}{param_name}="
underline = (" " * len(assign)) + ("^" * len(val))
msg = f"{msg}\n{assign}{val}\n{underline}"
raise TypeError(msg)
class _DeferredIterable(Generic[_T]):
"""Store a callable producing an iterable to defer collection until we need it."""
def __init__(self, into_iter: Callable[[], Iterable[_T]], /) -> None:
self._into_iter: Callable[[], Iterable[_T]] = into_iter
def __iter__(self) -> Iterator[_T]:
yield from self._into_iter()
def to_tuple(self) -> tuple[_T, ...]:
# Collect and return as a `tuple`.
it = self._into_iter()
return it if isinstance(it, tuple) else tuple(it)
Reference in a new issue Copy permalink