polars-runtime-compat 1.34.0b2__cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

polars/schema.py

@@ -0,0 +1,265 @@
+from __future__ import annotations
+
+import contextlib
+import sys
+from collections import OrderedDict
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Literal, Union, overload
+
+from polars._typing import PythonDataType
+from polars._utils.unstable import unstable
+from polars.datatypes import DataType, DataTypeClass, is_polars_dtype
+from polars.datatypes._parse import parse_into_dtype
+from polars.exceptions import DuplicateError
+from polars.interchange.protocol import CompatLevel
+
+with contextlib.suppress(ImportError):  # Module not available when building docs
+    from polars._plr import (
+        init_polars_schema_from_arrow_c_schema,
+        polars_schema_field_from_arrow_c_schema,
+        polars_schema_to_pycapsule,
+    )
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from polars import DataFrame, LazyFrame
+    from polars._typing import ArrowSchemaExportable
+
+    if sys.version_info >= (3, 10):
+        from typing import TypeAlias
+    else:
+        from typing_extensions import TypeAlias
+
+if sys.version_info >= (3, 10):
+
+    def _required_init_args(tp: DataTypeClass) -> bool:
+        # note: this check is ~20% faster than the check for a
+        # custom "__init__", below, but is not available on py39
+        return bool(tp.__annotations__)
+else:
+
+    def _required_init_args(tp: DataTypeClass) -> bool:
+        # indicates override of the default __init__
+        # (eg: this type requires specific args)
+        return "__init__" in tp.__dict__
+
+
+BaseSchema = OrderedDict[str, DataType]
+SchemaInitDataType: TypeAlias = Union[DataType, DataTypeClass, PythonDataType]
+
+__all__ = ["Schema"]
+
+
+def _check_dtype(tp: DataType | DataTypeClass) -> DataType:
+    if not isinstance(tp, DataType):
+        # note: if nested/decimal, or has signature params, this implies required args
+        if tp.is_nested() or tp.is_decimal() or _required_init_args(tp):
+            msg = f"dtypes must be fully-specified, got: {tp!r}"
+            raise TypeError(msg)
+        tp = tp()
+    return tp  # type: ignore[return-value]
+
+
+class Schema(BaseSchema):
+    """
+    Ordered mapping of column names to their data type.
+
+    Parameters
+    ----------
+    schema
+        The schema definition given by column names and their associated
+        Polars data type. Accepts a mapping, or an iterable of tuples, or any
+        object implementing the  `__arrow_c_schema__` PyCapsule interface
+        (e.g. pyarrow schemas).
+
+    Examples
+    --------
+    Define a schema by passing instantiated data types.
+
+    >>> schema = pl.Schema(
+    ...     {
+    ...         "foo": pl.String(),
+    ...         "bar": pl.Duration("us"),
+    ...         "baz": pl.Array(pl.Int8, 4),
+    ...     }
+    ... )
+    >>> schema
+    Schema({'foo': String, 'bar': Duration(time_unit='us'), 'baz': Array(Int8, shape=(4,))})
+
+    Access the data type associated with a specific column name.
+
+    >>> schema["baz"]
+    Array(Int8, shape=(4,))
+
+    Access various schema properties using the `names`, `dtypes`, and `len` methods.
+
+    >>> schema.names()
+    ['foo', 'bar', 'baz']
+    >>> schema.dtypes()
+    [String, Duration(time_unit='us'), Array(Int8, shape=(4,))]
+    >>> schema.len()
+    3
+
+    Import a pyarrow schema.
+
+    >>> import pyarrow as pa
+    >>> pl.Schema(pa.schema([pa.field("x", pa.int32())]))
+    Schema({'x': Int32})
+
+    Export a schema to pyarrow.
+
+    >>> pa.schema(pl.Schema({"x": pl.Int32}))
+    x: int32
+    """  # noqa: W505
+
+    def __init__(
+        self,
+        schema: (
+            Mapping[str, SchemaInitDataType]
+            | Iterable[tuple[str, SchemaInitDataType] | ArrowSchemaExportable]
+            | ArrowSchemaExportable
+            | None
+        ) = None,
+        *,
+        check_dtypes: bool = True,
+    ) -> None:
+        if hasattr(schema, "__arrow_c_schema__") and not isinstance(schema, Schema):
+            init_polars_schema_from_arrow_c_schema(self, schema)
+            return
+
+        input = schema.items() if isinstance(schema, Mapping) else (schema or ())
+        for v in input:
+            name, tp = (
+                polars_schema_field_from_arrow_c_schema(v)
+                if hasattr(v, "__arrow_c_schema__") and not isinstance(v, DataType)
+                else v
+            )
+
+            if name in self:
+                msg = f"iterable passed to pl.Schema contained duplicate name '{name}'"
+                raise DuplicateError(msg)
+
+            if not check_dtypes:
+                super().__setitem__(name, tp)  # type: ignore[assignment]
+            elif is_polars_dtype(tp):
+                super().__setitem__(name, _check_dtype(tp))
+            else:
+                self[name] = tp
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Mapping):
+            return False
+        if len(self) != len(other):
+            return False
+        for (nm1, tp1), (nm2, tp2) in zip(self.items(), other.items()):
+            if nm1 != nm2 or not tp1.is_(tp2):
+                return False
+        return True
+
+    def __ne__(self, other: object) -> bool:
+        return not self.__eq__(other)
+
+    def __setitem__(
+        self, name: str, dtype: DataType | DataTypeClass | PythonDataType
+    ) -> None:
+        dtype = _check_dtype(parse_into_dtype(dtype))
+        super().__setitem__(name, dtype)
+
+    @unstable()
+    def __arrow_c_schema__(self) -> object:
+        """
+        Export a Schema via the Arrow PyCapsule Interface.
+
+        https://arrow.apache.org/docs/dev/format/CDataInterface/PyCapsuleInterface.html
+        """
+        return polars_schema_to_pycapsule(self, CompatLevel.newest()._version)
+
+    def names(self) -> list[str]:
+        """
+        Get the column names of the schema.
+
+        Examples
+        --------
+        >>> s = pl.Schema({"x": pl.Float64(), "y": pl.Datetime(time_zone="UTC")})
+        >>> s.names()
+        ['x', 'y']
+        """
+        return list(self.keys())
+
+    def dtypes(self) -> list[DataType]:
+        """
+        Get the data types of the schema.
+
+        Examples
+        --------
+        >>> s = pl.Schema({"x": pl.UInt8(), "y": pl.List(pl.UInt8)})
+        >>> s.dtypes()
+        [UInt8, List(UInt8)]
+        """
+        return list(self.values())
+
+    @overload
+    def to_frame(self, *, eager: Literal[False]) -> LazyFrame: ...
+
+    @overload
+    def to_frame(self, *, eager: Literal[True] = ...) -> DataFrame: ...
+
+    def to_frame(self, *, eager: bool = True) -> DataFrame | LazyFrame:
+        """
+        Create an empty DataFrame (or LazyFrame) from this Schema.
+
+        Parameters
+        ----------
+        eager
+            If True, create a DataFrame; otherwise, create a LazyFrame.
+
+        Examples
+        --------
+        >>> s = pl.Schema({"x": pl.Int32(), "y": pl.String()})
+        >>> s.to_frame()
+        shape: (0, 2)
+        ┌─────┬─────┐
+        │ x   ┆ y   │
+        │ --- ┆ --- │
+        │ i32 ┆ str │
+        ╞═════╪═════╡
+        └─────┴─────┘
+        >>> s.to_frame(eager=False)  # doctest: +IGNORE_RESULT
+        <LazyFrame at 0x11BC0AD80>
+        """
+        from polars import DataFrame, LazyFrame
+
+        return DataFrame(schema=self) if eager else LazyFrame(schema=self)
+
+    def len(self) -> int:
+        """
+        Get the number of schema entries.
+
+        Examples
+        --------
+        >>> s = pl.Schema({"x": pl.Int32(), "y": pl.List(pl.String)})
+        >>> s.len()
+        2
+        >>> len(s)
+        2
+        """
+        return len(self)
+
+    def to_python(self) -> dict[str, type]:
+        """
+        Return a dictionary of column names and Python types.
+
+        Examples
+        --------
+        >>> s = pl.Schema(
+        ...     {
+        ...         "x": pl.Int8(),
+        ...         "y": pl.String(),
+        ...         "z": pl.Duration("us"),
+        ...     }
+        ... )
+        >>> s.to_python()
+        {'x': <class 'int'>, 'y':  <class 'str'>, 'z': <class 'datetime.timedelta'>}
+        """
+        return {name: tp.to_python() for name, tp in self.items()}