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

polars/datatype_expr/struct.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import polars._reexport as pl
+from polars._utils.various import qualified_type_name
+
+
+class DataTypeExprStructNameSpace:
+    """Namespace for struct datatype expressions."""
+
+    _accessor = "struct"
+
+    def __init__(self, expr: pl.DataTypeExpr) -> None:
+        self._pydatatype_expr = expr._pydatatype_expr
+
+    def __getitem__(self, item: str | int) -> pl.DataTypeExpr:
+        if isinstance(item, str):
+            return self.field_dtype(item)
+        elif isinstance(item, int):
+            return pl.DataTypeExpr._from_pydatatype_expr(
+                self._pydatatype_expr.struct_field_dtype_by_index(item)
+            )
+        else:
+            msg = f"expected type 'int | str', got {qualified_type_name(item)!r} ({item!r})"
+            raise TypeError(msg)
+
+    def field_dtype(self, field_name: str) -> pl.DataTypeExpr:
+        """
+
+        Get the DataType of field with a specific field name.
+
+        Notes
+        -----
+        The `struct` namespace has implemented `__getitem__` so you can also access
+        fields by index:
+
+        >>> (
+        ...     pl.Struct({"x": pl.Int64, "y": pl.String})
+        ...     .to_dtype_expr()
+        ...     .struct[1]
+        ...     .collect_dtype({})
+        ... )
+        String
+        """
+        return pl.DataTypeExpr._from_pydatatype_expr(
+            self._pydatatype_expr.struct_field_dtype_by_name(field_name)
+        )
+
+    def field_names(self) -> pl.Expr:
+        """
+        Get the field names in a struct as a list.
+
+        Examples
+        --------
+        >>> pl.select(
+        ...     pl.Struct({"x": pl.Int64, "y": pl.String})
+        ...     .to_dtype_expr()
+        ...     .struct.field_names()
+        ... )
+        shape: (2, 1)
+        ┌─────────┐
+        │ literal │
+        │ ---     │
+        │ str     │
+        ╞═════════╡
+        │ x       │
+        │ y       │
+        └─────────┘
+        """
+        return pl.Expr._from_pyexpr(self._pydatatype_expr.struct_field_names())

polars/datatypes/__init__.py

@@ -0,0 +1,122 @@
+from polars.datatypes._parse import (
+    parse_into_datatype_expr,
+    parse_into_dtype,
+    try_parse_into_dtype,
+)
+from polars.datatypes.classes import (
+    Array,
+    Binary,
+    Boolean,
+    Categorical,
+    Categories,
+    DataType,
+    DataTypeClass,
+    Date,
+    Datetime,
+    Decimal,
+    Duration,
+    Enum,
+    Field,
+    Float32,
+    Float64,
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    Int128,
+    IntegerType,
+    List,
+    Null,
+    Object,
+    String,
+    Struct,
+    TemporalType,
+    Time,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+    UInt128,
+    Unknown,
+    Utf8,
+)
+from polars.datatypes.constants import (
+    DTYPE_TEMPORAL_UNITS,
+    N_INFER_DEFAULT,
+)
+from polars.datatypes.constructor import (
+    numpy_type_to_constructor,
+    numpy_values_and_dtype,
+    polars_type_to_constructor,
+    py_type_to_constructor,
+)
+from polars.datatypes.convert import (
+    dtype_to_ffiname,
+    dtype_to_py_type,
+    is_polars_dtype,
+    maybe_cast,
+    numpy_char_code_to_dtype,
+    py_type_to_arrow_type,
+    supported_numpy_char_code,
+    unpack_dtypes,
+)
+
+__all__ = [
+    # classes
+    "Array",
+    "Binary",
+    "Boolean",
+    "Categorical",
+    "Categories",
+    "DataType",
+    "DataTypeClass",
+    "Date",
+    "Datetime",
+    "Decimal",
+    "Duration",
+    "Enum",
+    "Field",
+    "Float32",
+    "Float64",
+    "Int16",
+    "Int128",
+    "Int32",
+    "Int64",
+    "Int8",
+    "IntegerType",
+    "List",
+    "Null",
+    "Object",
+    "String",
+    "Struct",
+    "TemporalType",
+    "Time",
+    "UInt16",
+    "UInt128",
+    "UInt32",
+    "UInt64",
+    "UInt8",
+    "Unknown",
+    "Utf8",
+    # constants
+    "N_INFER_DEFAULT",
+    "DTYPE_TEMPORAL_UNITS",
+    # constructor
+    "numpy_type_to_constructor",
+    "numpy_values_and_dtype",
+    "polars_type_to_constructor",
+    "py_type_to_constructor",
+    # convert
+    "dtype_to_ffiname",
+    "dtype_to_py_type",
+    "is_polars_dtype",
+    "maybe_cast",
+    "numpy_char_code_to_dtype",
+    "py_type_to_arrow_type",
+    "supported_numpy_char_code",
+    "unpack_dtypes",
+    # _parse
+    "parse_into_dtype",
+    "parse_into_datatype_expr",
+    "try_parse_into_dtype",
+]

polars/datatypes/_parse.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import enum
+import functools
+import re
+import sys
+from datetime import date, datetime, time, timedelta
+from decimal import Decimal as PyDecimal
+from inspect import isclass
+from typing import TYPE_CHECKING, Any, ForwardRef, NoReturn, Union, get_args
+
+import polars._reexport as pl
+from polars.datatypes.classes import (
+    Binary,
+    Boolean,
+    Date,
+    Datetime,
+    Decimal,
+    Duration,
+    Enum,
+    Float64,
+    Int64,
+    List,
+    Null,
+    Object,
+    String,
+    Time,
+    Unknown,
+)
+from polars.datatypes.convert import is_polars_dtype
+
+if TYPE_CHECKING:
+    from polars._typing import PolarsDataType, PythonDataType, SchemaDict
+
+
+UnionTypeOld = type(Union[int, str])
+if sys.version_info >= (3, 10):
+    from types import NoneType, UnionType
+else:  # pragma: no cover
+    # Define equivalent for older Python versions
+    NoneType = type(None)
+    UnionType = UnionTypeOld
+
+
+def parse_into_datatype_expr(input: Any) -> pl.DataTypeExpr:
+    """Parse an input into a DataTypeExpr."""
+    if isinstance(input, pl.DataTypeExpr):
+        return input
+    else:
+        return parse_into_dtype(input).to_dtype_expr()
+
+
+def parse_into_dtype(input: Any) -> PolarsDataType:
+    """
+    Parse an input into a Polars data type.
+
+    Raises
+    ------
+    TypeError
+        If the input cannot be parsed into a Polars data type.
+    """
+    if is_polars_dtype(input):
+        return input
+    elif isinstance(input, ForwardRef):
+        return _parse_forward_ref_into_dtype(input)
+    elif isinstance(input, (UnionType, UnionTypeOld)):
+        return _parse_union_type_into_dtype(input)
+    else:
+        return parse_py_type_into_dtype(input)
+
+
+def try_parse_into_dtype(input: Any) -> PolarsDataType | None:
+    """Try parsing an input into a Polars data type, returning None on failure."""
+    try:
+        return parse_into_dtype(input)
+    except TypeError:
+        return None
+
+
+@functools.lru_cache(16)
+def parse_py_type_into_dtype(input: PythonDataType | type[object]) -> PolarsDataType:
+    """Convert Python data type to Polars data type."""
+    if input is int:
+        return Int64()
+    elif input is float:
+        return Float64()
+    elif input is str:
+        return String()
+    elif input is bool:
+        return Boolean()
+    elif isinstance(input, type) and issubclass(input, datetime):  # type: ignore[redundant-expr]
+        return Datetime("us")
+    elif isinstance(input, type) and issubclass(input, date):  # type: ignore[redundant-expr]
+        return Date()
+    elif input is timedelta:
+        return Duration
+    elif input is time:
+        return Time()
+    elif input is PyDecimal:
+        return Decimal
+    elif input is bytes:
+        return Binary()
+    elif input is object:
+        return Object()
+    elif input is NoneType:
+        return Null()
+    elif input is list or input is tuple:
+        return List
+    elif isclass(input) and issubclass(input, enum.Enum):
+        return Enum(input)
+    # this is required as pass through. Don't remove
+    elif input == Unknown:
+        return Unknown
+    elif hasattr(input, "__origin__") and hasattr(input, "__args__"):
+        return _parse_generic_into_dtype(input)
+    else:
+        _raise_on_invalid_dtype(input)
+
+
+def _parse_generic_into_dtype(input: Any) -> PolarsDataType:
+    """Parse a generic type (from typing annotation) into a Polars data type."""
+    base_type = input.__origin__
+    if base_type not in (tuple, list):
+        _raise_on_invalid_dtype(input)
+
+    inner_types = input.__args__
+    inner_type = inner_types[0]
+    if len(inner_types) > 1:
+        all_equal = all(t in (inner_type, ...) for t in inner_types)
+        if not all_equal:
+            _raise_on_invalid_dtype(input)
+
+    inner_type = inner_types[0]
+    inner_dtype = parse_py_type_into_dtype(inner_type)
+    return List(inner_dtype)
+
+
+PY_TYPE_STR_TO_DTYPE: SchemaDict = {
+    "Decimal": Decimal,
+    "NoneType": Null(),
+    "bool": Boolean(),
+    "bytes": Binary(),
+    "date": Date(),
+    "datetime": Datetime("us"),
+    "float": Float64(),
+    "int": Int64(),
+    "list": List,
+    "object": Object(),
+    "str": String(),
+    "time": Time(),
+    "timedelta": Duration,
+    "tuple": List,
+}
+
+
+def _parse_forward_ref_into_dtype(input: ForwardRef) -> PolarsDataType:
+    """Parse a ForwardRef into a Polars data type."""
+    annotation = input.__forward_arg__
+
+    # Strip "optional" designation - Polars data types are always nullable
+    formatted = re.sub(r"(^None \|)|(\| None$)", "", annotation).strip()
+
+    try:
+        return PY_TYPE_STR_TO_DTYPE[formatted]
+    except KeyError:
+        _raise_on_invalid_dtype(input)
+
+
+def _parse_union_type_into_dtype(input: Any) -> PolarsDataType:
+    """
+    Parse a union of types into a Polars data type.
+
+    Unions of multiple non-null types (e.g. `int | float`) are not supported.
+
+    Parameters
+    ----------
+    input
+        A union type, e.g. `str | None` (new syntax) or `Union[str, None]` (old syntax).
+    """
+    # Strip "optional" designation - Polars data types are always nullable
+    inner_types = [tp for tp in get_args(input) if tp is not NoneType]
+
+    if len(inner_types) != 1:
+        _raise_on_invalid_dtype(input)
+
+    input = inner_types[0]
+    return parse_into_dtype(input)
+
+
+def _raise_on_invalid_dtype(input: Any) -> NoReturn:
+    """Raise an informative error if the input could not be parsed."""
+    input_type = input if type(input) is type else f"of type {type(input).__name__!r}"
+    input_detail = "" if type(input) is type else f" (given: {input!r})"
+    msg = f"cannot parse input {input_type} into Polars data type{input_detail}"
+    raise TypeError(msg) from None

polars/datatypes/_utils.py

@@ -0,0 +1,48 @@
+"""Utility functions for handling and processing of datatypes."""
+
+from polars._typing import PolarsDataType
+from polars.datatypes.classes import Array, List, Struct
+
+
+def dtype_to_init_repr(dtype: PolarsDataType, prefix: str = "pl.") -> str:
+    """Convert a Polars dtype to a prefixed string representation."""
+    if isinstance(dtype, List):
+        init_repr = _dtype_to_init_repr_list(dtype, prefix)
+    elif isinstance(dtype, Array):
+        init_repr = _dtype_to_init_repr_array(dtype, prefix)
+    elif isinstance(dtype, Struct):
+        init_repr = _dtype_to_init_repr_struct(dtype, prefix)
+    else:
+        init_repr = f"{prefix}{dtype!r}"
+    return init_repr
+
+
+def _dtype_to_init_repr_list(dtype: List, prefix: str) -> str:
+    class_name = dtype.__class__.__name__
+    if dtype.inner is not None:
+        inner_repr = dtype_to_init_repr(dtype.inner, prefix)
+    else:
+        inner_repr = ""
+    init_repr = f"{prefix}{class_name}({inner_repr})"
+    return init_repr
+
+
+def _dtype_to_init_repr_array(dtype: Array, prefix: str) -> str:
+    class_name = dtype.__class__.__name__
+    if dtype.inner is not None:
+        inner_repr = dtype_to_init_repr(dtype.inner, prefix)
+    else:
+        inner_repr = ""
+    init_repr = f"{prefix}{class_name}({inner_repr}, shape={dtype.shape})"
+    return init_repr
+
+
+def _dtype_to_init_repr_struct(dtype: Struct, prefix: str) -> str:
+    class_name = dtype.__class__.__name__
+    inner_list = [
+        f"{field_name!r}: {dtype_to_init_repr(inner_dtype, prefix)}"
+        for field_name, inner_dtype in dict(dtype).items()
+    ]
+    inner_repr = "{" + ", ".join(inner_list) + "}"
+    init_repr = f"{prefix}{class_name}({inner_repr})"
+    return init_repr