polars-runtime-compat 1.34.0b2__cp39-abi3-win_amd64.whl

polars/datatypes/constants.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from polars._typing import TimeUnit
+
+# Number of rows to scan by default when inferring datatypes
+N_INFER_DEFAULT = 100
+
+DTYPE_TEMPORAL_UNITS: frozenset[TimeUnit] = frozenset(["ns", "us", "ms"])

polars/datatypes/constructor.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import functools
+from decimal import Decimal as PyDecimal
+from typing import TYPE_CHECKING, Any, Callable
+
+from polars import datatypes as dt
+from polars._dependencies import numpy as np
+
+# Module not available when building docs
+try:
+    from polars._plr import PySeries
+
+    _DOCUMENTING = False
+except ImportError:
+    _DOCUMENTING = True
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from polars._typing import PolarsDataType
+
+if not _DOCUMENTING:
+    _POLARS_TYPE_TO_CONSTRUCTOR: dict[
+        PolarsDataType, Callable[[str, Sequence[Any], bool], PySeries]
+    ] = {
+        dt.Float32: PySeries.new_opt_f32,
+        dt.Float64: PySeries.new_opt_f64,
+        dt.Int8: PySeries.new_opt_i8,
+        dt.Int16: PySeries.new_opt_i16,
+        dt.Int32: PySeries.new_opt_i32,
+        dt.Int64: PySeries.new_opt_i64,
+        dt.Int128: PySeries.new_opt_i128,
+        dt.UInt8: PySeries.new_opt_u8,
+        dt.UInt16: PySeries.new_opt_u16,
+        dt.UInt32: PySeries.new_opt_u32,
+        dt.UInt64: PySeries.new_opt_u64,
+        dt.UInt128: PySeries.new_opt_u128,
+        dt.Decimal: PySeries.new_decimal,
+        dt.Date: PySeries.new_opt_i32,
+        dt.Datetime: PySeries.new_opt_i64,
+        dt.Duration: PySeries.new_opt_i64,
+        dt.Time: PySeries.new_opt_i64,
+        dt.Boolean: PySeries.new_opt_bool,
+        dt.String: PySeries.new_str,
+        dt.Object: PySeries.new_object,
+        dt.Categorical: PySeries.new_str,
+        dt.Enum: PySeries.new_str,
+        dt.Binary: PySeries.new_binary,
+        dt.Null: PySeries.new_null,
+    }
+
+
+def polars_type_to_constructor(
+    dtype: PolarsDataType,
+) -> Callable[[str, Sequence[Any], bool], PySeries]:
+    """Get the right PySeries constructor for the given Polars dtype."""
+    # Special case for Array as it needs to pass the dtype argument on construction
+    if isinstance(dtype, dt.Array):
+        return functools.partial(PySeries.new_array, dtype=dtype)
+
+    try:
+        base_type = dtype.base_type()
+        return _POLARS_TYPE_TO_CONSTRUCTOR[base_type]
+    except KeyError:  # pragma: no cover
+        msg = f"cannot construct PySeries for type {dtype!r}"
+        raise ValueError(msg) from None
+
+
+_NUMPY_TYPE_TO_CONSTRUCTOR = None
+
+
+def _set_numpy_to_constructor() -> None:
+    global _NUMPY_TYPE_TO_CONSTRUCTOR
+    _NUMPY_TYPE_TO_CONSTRUCTOR = {
+        np.float32: PySeries.new_f32,
+        np.float64: PySeries.new_f64,
+        np.int8: PySeries.new_i8,
+        np.int16: PySeries.new_i16,
+        np.int32: PySeries.new_i32,
+        np.int64: PySeries.new_i64,
+        np.uint8: PySeries.new_u8,
+        np.uint16: PySeries.new_u16,
+        np.uint32: PySeries.new_u32,
+        np.uint64: PySeries.new_u64,
+        np.str_: PySeries.new_str,
+        np.bytes_: PySeries.new_binary,
+        np.bool_: PySeries.new_bool,
+        np.datetime64: PySeries.new_i64,
+        np.timedelta64: PySeries.new_i64,
+    }
+
+
+@functools.lru_cache(maxsize=32)
+def _normalise_numpy_dtype(dtype: Any) -> tuple[Any, Any]:
+    normalised_dtype = (
+        np.dtype(dtype.base.name) if dtype.kind in ("i", "u", "f") else dtype
+    ).type
+    cast_as: Any = None
+    if normalised_dtype == np.float16:
+        normalised_dtype = cast_as = np.float32
+    elif normalised_dtype in (np.datetime64, np.timedelta64):
+        time_unit = np.datetime_data(dtype)[0]
+        if time_unit in dt.DTYPE_TEMPORAL_UNITS or (
+            time_unit == "D" and normalised_dtype == np.datetime64
+        ):
+            cast_as = np.int64
+        else:
+            msg = (
+                "incorrect NumPy datetime resolution"
+                "\n\n'D' (datetime only), 'ms', 'us', and 'ns' resolutions are supported when converting from numpy.{datetime64,timedelta64}."
+                " Please cast to the closest supported unit before converting."
+            )
+            raise ValueError(msg)
+    return normalised_dtype, cast_as
+
+
+def numpy_values_and_dtype(
+    values: np.ndarray[Any, Any],
+) -> tuple[np.ndarray[Any, Any], type]:
+    """Return numpy values and their associated dtype, adjusting if required."""
+    # Create new dtype object from dtype base name so architecture specific
+    # dtypes (np.longlong np.ulonglong np.intc np.uintc np.longdouble, ...)
+    # get converted to their normalized dtype (np.int*, np.uint*, np.float*).
+    dtype, cast_as = _normalise_numpy_dtype(values.dtype)
+    if cast_as:
+        values = values.astype(cast_as)
+    return values, dtype
+
+
+def numpy_type_to_constructor(
+    values: np.ndarray[Any, Any], dtype: type[np.dtype[Any]]
+) -> Callable[..., PySeries]:
+    """Get the right PySeries constructor for the given Polars dtype."""
+    if _NUMPY_TYPE_TO_CONSTRUCTOR is None:
+        _set_numpy_to_constructor()
+    try:
+        return _NUMPY_TYPE_TO_CONSTRUCTOR[dtype]  # type:ignore[index]
+    except KeyError:
+        if len(values) > 0:
+            first_non_nan = next(
+                (v for v in values if isinstance(v, np.ndarray) or v == v), None
+            )
+            if isinstance(first_non_nan, str):
+                return PySeries.new_str
+            if isinstance(first_non_nan, bytes):
+                return PySeries.new_binary
+        return PySeries.new_object
+    except NameError:  # pragma: no cover
+        msg = f"'numpy' is required to convert numpy dtype {dtype!r}"
+        raise ModuleNotFoundError(msg) from None
+
+
+if not _DOCUMENTING:
+    _PY_TYPE_TO_CONSTRUCTOR = {
+        float: PySeries.new_opt_f64,
+        bool: PySeries.new_opt_bool,
+        int: PySeries.new_opt_i64,
+        str: PySeries.new_str,
+        bytes: PySeries.new_binary,
+        PyDecimal: PySeries.new_decimal,
+    }
+
+
+def py_type_to_constructor(py_type: type[Any]) -> Callable[..., PySeries]:
+    """Get the right PySeries constructor for the given Python dtype."""
+    py_type = (
+        next((tp for tp in _PY_TYPE_TO_CONSTRUCTOR if issubclass(py_type, tp)), py_type)
+        if py_type not in _PY_TYPE_TO_CONSTRUCTOR
+        else py_type
+    )
+    return _PY_TYPE_TO_CONSTRUCTOR.get(py_type, PySeries.new_object)

polars/datatypes/convert.py

@@ -0,0 +1,366 @@
+from __future__ import annotations
+
+import contextlib
+import functools
+import re
+import sys
+from collections.abc import Collection
+from datetime import date, datetime, time, timedelta
+from decimal import Decimal as PyDecimal
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+from polars._dependencies import numpy as np
+from polars._dependencies import pyarrow as pa
+from polars.datatypes.classes import (
+    Array,
+    Binary,
+    Boolean,
+    Categorical,
+    DataType,
+    DataTypeClass,
+    Date,
+    Datetime,
+    Decimal,
+    Duration,
+    Enum,
+    Field,
+    Float32,
+    Float64,
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    Int128,
+    List,
+    Null,
+    Object,
+    String,
+    Struct,
+    Time,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+    UInt128,
+    Unknown,
+)
+
+with contextlib.suppress(ImportError):  # Module not available when building docs
+    from polars._plr import dtype_str_repr as _dtype_str_repr
+
+
+OptionType = type(Optional[type])
+if sys.version_info >= (3, 10):
+    from types import NoneType, UnionType
+else:
+    # infer equivalent class
+    NoneType = type(None)
+    UnionType = type(Union[int, float])
+
+if TYPE_CHECKING:
+    from polars._typing import PolarsDataType, PythonDataType, TimeUnit
+
+    if sys.version_info >= (3, 10):
+        from typing import TypeGuard
+    else:
+        from typing_extensions import TypeGuard
+
+
+def is_polars_dtype(
+    dtype: Any,
+    *,
+    include_unknown: bool = False,
+    require_instantiated: bool = False,
+) -> TypeGuard[PolarsDataType]:
+    """Indicate whether the given input is a Polars dtype, or dtype specialization."""
+    check_classes = DataType if require_instantiated else (DataType, DataTypeClass)
+    is_dtype = isinstance(dtype, check_classes)
+
+    if not include_unknown:
+        return is_dtype and dtype != Unknown
+    else:
+        return is_dtype
+
+
+def unpack_dtypes(
+    *dtypes: PolarsDataType | None,
+    include_compound: bool = False,
+) -> set[PolarsDataType]:
+    """
+    Return a set of unique dtypes found in one or more (potentially compound) dtypes.
+
+    Parameters
+    ----------
+    *dtypes
+        One or more Polars dtypes.
+    include_compound
+        * if True, any parent/compound dtypes (List, Struct) are included in the result.
+        * if False, only the child/scalar dtypes are returned from these types.
+
+    Examples
+    --------
+    >>> from polars.datatypes import unpack_dtypes
+    >>> list_dtype = [pl.List(pl.Float64)]
+    >>> struct_dtype = pl.Struct(
+    ...     [
+    ...         pl.Field("a", pl.Int64),
+    ...         pl.Field("b", pl.String),
+    ...         pl.Field("c", pl.List(pl.Float64)),
+    ...     ]
+    ... )
+    >>> unpack_dtypes([struct_dtype, list_dtype])  # doctest: +IGNORE_RESULT
+    {Float64, Int64, String}
+    >>> unpack_dtypes(
+    ...     [struct_dtype, list_dtype], include_compound=True
+    ... )  # doctest: +IGNORE_RESULT
+    {Float64, Int64, String, List(Float64), Struct([Field('a', Int64), Field('b', String), Field('c', List(Float64))])}
+    """  # noqa: W505
+    if not dtypes:
+        return set()
+    elif len(dtypes) == 1 and isinstance(dtypes[0], Collection):
+        dtypes = dtypes[0]
+
+    unpacked: set[PolarsDataType] = set()
+    for tp in dtypes:
+        if isinstance(tp, (List, Array)):
+            if include_compound:
+                unpacked.add(tp)
+            unpacked.update(unpack_dtypes(tp.inner, include_compound=include_compound))
+        elif isinstance(tp, Struct):
+            if include_compound:
+                unpacked.add(tp)
+            unpacked.update(unpack_dtypes(tp.fields, include_compound=include_compound))  # type: ignore[arg-type]
+        elif isinstance(tp, Field):
+            unpacked.update(unpack_dtypes(tp.dtype, include_compound=include_compound))
+        elif tp is not None and is_polars_dtype(tp):
+            unpacked.add(tp)
+    return unpacked
+
+
+class _DataTypeMappings:
+    @property
+    @functools.lru_cache  # noqa: B019
+    def DTYPE_TO_FFINAME(self) -> dict[PolarsDataType, str]:
+        return {
+            Binary: "binary",
+            Boolean: "bool",
+            Categorical: "categorical",
+            Date: "date",
+            Datetime: "datetime",
+            Decimal: "decimal",
+            Duration: "duration",
+            Float32: "f32",
+            Float64: "f64",
+            Int8: "i8",
+            Int16: "i16",
+            Int32: "i32",
+            Int64: "i64",
+            Int128: "i128",
+            List: "list",
+            Object: "object",
+            String: "str",
+            Struct: "struct",
+            Time: "time",
+            UInt8: "u8",
+            UInt16: "u16",
+            UInt32: "u32",
+            UInt64: "u64",
+            UInt128: "u128",
+        }
+
+    @property
+    @functools.lru_cache  # noqa: B019
+    def DTYPE_TO_PY_TYPE(self) -> dict[PolarsDataType, PythonDataType]:
+        return {
+            Array: list,
+            Binary: bytes,
+            Boolean: bool,
+            Date: date,
+            Datetime: datetime,
+            Decimal: PyDecimal,
+            Duration: timedelta,
+            Float32: float,
+            Float64: float,
+            Int8: int,
+            Int16: int,
+            Int32: int,
+            Int64: int,
+            Int128: int,
+            List: list,
+            Null: None.__class__,
+            Object: object,
+            String: str,
+            Struct: dict,
+            Time: time,
+            UInt8: int,
+            UInt16: int,
+            UInt32: int,
+            UInt64: int,
+            UInt128: int,
+            # the below mappings are appropriate as we restrict cat/enum to strings
+            Enum: str,
+            Categorical: str,
+        }
+
+    @property
+    @functools.lru_cache  # noqa: B019
+    def NUMPY_KIND_AND_ITEMSIZE_TO_DTYPE(self) -> dict[tuple[str, int], PolarsDataType]:
+        return {
+            # (np.dtype().kind, np.dtype().itemsize)
+            ("M", 8): Datetime,
+            ("b", 1): Boolean,
+            ("f", 2): Float32,
+            ("f", 4): Float32,
+            ("f", 8): Float64,
+            ("i", 1): Int8,
+            ("i", 2): Int16,
+            ("i", 4): Int32,
+            ("i", 8): Int64,
+            ("m", 8): Duration,
+            ("u", 1): UInt8,
+            ("u", 2): UInt16,
+            ("u", 4): UInt32,
+            ("u", 8): UInt64,
+        }
+
+    @property
+    @functools.lru_cache  # noqa: B019
+    def PY_TYPE_TO_ARROW_TYPE(self) -> dict[PythonDataType, pa.lib.DataType]:
+        return {
+            bool: pa.bool_(),
+            date: pa.date32(),
+            datetime: pa.timestamp("us"),
+            float: pa.float64(),
+            int: pa.int64(),
+            str: pa.large_utf8(),
+            time: pa.time64("us"),
+            timedelta: pa.duration("us"),
+            None.__class__: pa.null(),
+        }
+
+    @property
+    @functools.lru_cache  # noqa: B019
+    def REPR_TO_DTYPE(self) -> dict[str, PolarsDataType]:
+        def _dtype_str_repr_safe(o: Any) -> PolarsDataType | None:
+            try:
+                return _dtype_str_repr(o.base_type()).split("[")[0]  # type: ignore[return-value]
+            except TypeError:
+                return None
+
+        return {
+            _dtype_str_repr_safe(obj): obj  # type: ignore[misc]
+            for obj in globals().values()
+            if is_polars_dtype(obj) and _dtype_str_repr_safe(obj) is not None
+        }
+
+
+# Initialize once (poor man's singleton :)
+DataTypeMappings = _DataTypeMappings()
+
+
+def dtype_to_ffiname(dtype: PolarsDataType) -> str:
+    """Return FFI function name associated with the given Polars dtype."""
+    try:
+        dtype = dtype.base_type()
+        return DataTypeMappings.DTYPE_TO_FFINAME[dtype]
+    except KeyError:  # pragma: no cover
+        msg = f"conversion of polars data type {dtype!r} to FFI not implemented"
+        raise NotImplementedError(msg) from None
+
+
+def dtype_to_py_type(dtype: PolarsDataType) -> PythonDataType:
+    """Convert a Polars dtype to a Python dtype."""
+    try:
+        dtype = dtype.base_type()
+        return DataTypeMappings.DTYPE_TO_PY_TYPE[dtype]
+    except KeyError:  # pragma: no cover
+        msg = f"conversion of polars data type {dtype!r} to Python type not implemented"
+        raise NotImplementedError(msg) from None
+
+
+def py_type_to_arrow_type(dtype: PythonDataType) -> pa.lib.DataType:
+    """Convert a Python dtype to an Arrow dtype."""
+    try:
+        return DataTypeMappings.PY_TYPE_TO_ARROW_TYPE[dtype]
+    except KeyError:  # pragma: no cover
+        msg = f"cannot parse Python data type {dtype!r} into Arrow data type"
+        raise ValueError(msg) from None
+
+
+def dtype_short_repr_to_dtype(dtype_string: str | None) -> PolarsDataType | None:
+    """Map a PolarsDataType short repr (eg: 'i64', 'list[str]') back into a dtype."""
+    if dtype_string is None:
+        return None
+
+    m = re.match(r"^(\w+)(?:\[(.+)\])?$", dtype_string)
+    if m is None:
+        return None
+
+    dtype_base, subtype = m.groups()
+    dtype = DataTypeMappings.REPR_TO_DTYPE.get(dtype_base)
+    if dtype and subtype:
+        # TODO: further-improve handling for nested types (such as List,Struct)
+        try:
+            if dtype == Decimal:
+                subtype = (None, int(subtype))
+            else:
+                subtype = (
+                    s.strip("'\" ") for s in subtype.replace("μs", "us").split(",")
+                )
+            return dtype(*subtype)  # type: ignore[operator]
+        except ValueError:
+            pass
+    return dtype
+
+
+def supported_numpy_char_code(dtype_char: str) -> bool:
+    """Check if the input can be mapped to a Polars dtype."""
+    dtype = np.dtype(dtype_char)
+    return (
+        dtype.kind,
+        dtype.itemsize,
+    ) in DataTypeMappings.NUMPY_KIND_AND_ITEMSIZE_TO_DTYPE
+
+
+def numpy_char_code_to_dtype(dtype_char: str) -> PolarsDataType:
+    """Convert a numpy character dtype to a Polars dtype."""
+    dtype = np.dtype(dtype_char)
+    if dtype.kind == "U":
+        return String
+    elif dtype.kind == "S":
+        return Binary
+    try:
+        return DataTypeMappings.NUMPY_KIND_AND_ITEMSIZE_TO_DTYPE[
+            dtype.kind, dtype.itemsize
+        ]
+    except KeyError:  # pragma: no cover
+        msg = f"cannot parse numpy data type {dtype!r} into Polars data type"
+        raise ValueError(msg) from None
+
+
+def maybe_cast(el: Any, dtype: PolarsDataType) -> Any:
+    """Try casting a value to a value that is valid for the given Polars dtype."""
+    # cast el if it doesn't match
+    from polars._utils.convert import (
+        datetime_to_int,
+        timedelta_to_int,
+    )
+
+    time_unit: TimeUnit
+    if isinstance(el, datetime):
+        time_unit = getattr(dtype, "time_unit", "us")
+        return datetime_to_int(el, time_unit)
+    elif isinstance(el, timedelta):
+        time_unit = getattr(dtype, "time_unit", "us")
+        return timedelta_to_int(el, time_unit)
+
+    py_type = dtype_to_py_type(dtype)
+    if not isinstance(el, py_type):
+        try:
+            el = py_type(el)  # type: ignore[call-arg]
+        except Exception:
+            from polars._utils.various import qualified_type_name
+
+            msg = f"cannot convert Python type {qualified_type_name(el)!r} to {dtype!r}"
+            raise TypeError(msg) from None
+    return el

polars/datatypes/group.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from polars.datatypes.classes import (
+    Array,
+    DataType,
+    DataTypeClass,
+    Date,
+    Datetime,
+    Decimal,
+    Duration,
+    Float32,
+    Float64,
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    Int128,
+    List,
+    Struct,
+    Time,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+    UInt128,
+)
+
+if TYPE_CHECKING:
+    import sys
+    from collections.abc import Iterable
+
+    from polars._typing import (
+        PolarsDataType,
+        PolarsIntegerType,
+        PolarsTemporalType,
+    )
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
+
+class DataTypeGroup(frozenset):  # type: ignore[type-arg]
+    """Group of data types."""
+
+    _match_base_type: bool
+
+    def __new__(
+        cls, items: Iterable[DataType | DataTypeClass], *, match_base_type: bool = True
+    ) -> Self:
+        """
+        Construct a DataTypeGroup.
+
+        Parameters
+        ----------
+        items :
+            iterable of data types
+        match_base_type:
+            match the base type
+        """
+        for it in items:
+            if not isinstance(it, (DataType, DataTypeClass)):
+                from polars._utils.various import qualified_type_name
+
+                msg = f"DataTypeGroup items must be dtypes; found {qualified_type_name(it)!r}"
+                raise TypeError(msg)
+
+        dtype_group = super().__new__(cls, items)
+        dtype_group._match_base_type = match_base_type
+        return dtype_group
+
+    def __contains__(self, item: Any) -> bool:
+        if self._match_base_type and isinstance(item, (DataType, DataTypeClass)):
+            item = item.base_type()
+        return super().__contains__(item)
+
+
+SIGNED_INTEGER_DTYPES: frozenset[PolarsIntegerType] = DataTypeGroup(
+    [
+        Int8,
+        Int16,
+        Int32,
+        Int64,
+        Int128,
+    ]
+)
+UNSIGNED_INTEGER_DTYPES: frozenset[PolarsIntegerType] = DataTypeGroup(
+    [
+        UInt8,
+        UInt16,
+        UInt32,
+        UInt64,
+        UInt128,
+    ]
+)
+INTEGER_DTYPES: frozenset[PolarsIntegerType] = (
+    SIGNED_INTEGER_DTYPES | UNSIGNED_INTEGER_DTYPES
+)
+FLOAT_DTYPES: frozenset[PolarsDataType] = DataTypeGroup([Float32, Float64])
+NUMERIC_DTYPES: frozenset[PolarsDataType] = DataTypeGroup(
+    FLOAT_DTYPES | INTEGER_DTYPES | frozenset([Decimal])
+)
+
+DATETIME_DTYPES: frozenset[PolarsDataType] = DataTypeGroup(
+    [
+        Datetime,
+        Datetime("ms"),
+        Datetime("us"),
+        Datetime("ns"),
+        Datetime("ms", "*"),
+        Datetime("us", "*"),
+        Datetime("ns", "*"),
+    ]
+)
+DURATION_DTYPES: frozenset[PolarsDataType] = DataTypeGroup(
+    [
+        Duration,
+        Duration("ms"),
+        Duration("us"),
+        Duration("ns"),
+    ]
+)
+TEMPORAL_DTYPES: frozenset[PolarsTemporalType] = DataTypeGroup(
+    frozenset([Date, Time]) | DATETIME_DTYPES | DURATION_DTYPES
+)
+
+NESTED_DTYPES: frozenset[PolarsDataType] = DataTypeGroup([List, Struct, Array])