diff --git a/dataframely/_pydantic.py b/dataframely/_pydantic.py index ddcd44d6..7f8113f2 100644 --- a/dataframely/_pydantic.py +++ b/dataframely/_pydantic.py @@ -18,10 +18,7 @@ def _dict_to_df(schema_type: type[BaseSchema], data: dict) -> pl.DataFrame: - return pl.from_dict( - data, - schema=schema_type.to_polars_schema(), # type: ignore[attr-defined] - ) + return pl.from_dict(data, schema=pl.Schema(schema_type)) def _validate_df_schema(schema_type: type[_S], df: pl.DataFrame) -> DataFrame[_S]: diff --git a/dataframely/filter_result.py b/dataframely/filter_result.py index 8e53e117..c7ad6fc1 100644 --- a/dataframely/filter_result.py +++ b/dataframely/filter_result.py @@ -111,7 +111,7 @@ def _create_empty(cls, schema: type[S], with_casting_rules: bool) -> FailureInfo rules = schema._validation_rules(with_cast=with_casting_rules) lf = pl.LazyFrame( schema={ - **schema.to_polars_schema(), # type: ignore + **pl.Schema(schema), **{rule: pl.Boolean for rule in rules}, } ) diff --git a/dataframely/schema.py b/dataframely/schema.py index 15602a3c..bc8dd336 100644 --- a/dataframely/schema.py +++ b/dataframely/schema.py @@ -1385,15 +1385,6 @@ def _validate_if_needed( # ----------------------------- THIRD-PARTY PACKAGES ----------------------------- # - @classmethod - def to_polars_schema(cls) -> pl.Schema: - """Obtain the polars schema for this schema. - - Returns: - A :mod:`polars` schema that mirrors the schema defined by this class. - """ - return pl.Schema({name: col.dtype for name, col in cls.columns().items()}) - @classmethod def to_sqlalchemy_columns(cls, dialect: sa.Dialect) -> list[sa.Column]: """Obtain the SQLAlchemy column definitions for a particular dialect for this diff --git a/docs/api/schema/conversion.rst b/docs/api/schema/conversion.rst index 074049c6..379efc95 100644 --- a/docs/api/schema/conversion.rst +++ b/docs/api/schema/conversion.rst @@ -7,4 +7,3 @@ Conversion :toctree: _gen/ Schema.to_sqlalchemy_columns - Schema.to_polars_schema diff --git a/tests/collection/test_base.py b/tests/collection/test_base.py index 4b045393..b6084278 100644 --- a/tests/collection/test_base.py +++ b/tests/collection/test_base.py @@ -56,9 +56,9 @@ def test_cast() -> None: "second": pl.LazyFrame({"a": [1, 2, 3], "b": [4, 5, 6]}), }, ) - assert collection.first.collect_schema() == MyFirstSchema.to_polars_schema() + assert collection.first.collect_schema() == pl.Schema(MyFirstSchema) assert collection.second is not None - assert collection.second.collect_schema() == MySecondSchema.to_polars_schema() + assert collection.second.collect_schema() == pl.Schema(MySecondSchema) @pytest.mark.parametrize( diff --git a/tests/collection/test_cast.py b/tests/collection/test_cast.py index 5893bc74..08d31838 100644 --- a/tests/collection/test_cast.py +++ b/tests/collection/test_cast.py @@ -26,16 +26,16 @@ def test_cast_valid(df_type: type[pl.DataFrame] | type[pl.LazyFrame]) -> None: first = df_type({"a": [3]}) second = df_type({"a": [1]}) out = Collection.cast({"first": first, "second": second}) # type: ignore - assert out.first.collect_schema() == FirstSchema.to_polars_schema() + assert out.first.collect_schema() == pl.Schema(FirstSchema) assert out.second is not None - assert out.second.collect_schema() == SecondSchema.to_polars_schema() + assert out.second.collect_schema() == pl.Schema(SecondSchema) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) def test_cast_valid_optional(df_type: type[pl.DataFrame] | type[pl.LazyFrame]) -> None: first = df_type({"a": [3]}) out = Collection.cast({"first": first}) # type: ignore - assert out.first.collect_schema() == FirstSchema.to_polars_schema() + assert out.first.collect_schema() == pl.Schema(FirstSchema) assert out.second is None diff --git a/tests/collection/test_create_empty.py b/tests/collection/test_create_empty.py index 60e5ca48..66f510f1 100644 --- a/tests/collection/test_create_empty.py +++ b/tests/collection/test_create_empty.py @@ -1,6 +1,7 @@ # Copyright (c) QuantCo 2025-2026 # SPDX-License-Identifier: BSD-3-Clause +import polars as pl import dataframely as dy @@ -23,7 +24,7 @@ class MyCollection(dy.Collection): def test_create_empty() -> None: collection = MyCollection.create_empty() assert collection.first.collect().height == 0 - assert collection.first.collect_schema() == MyFirstSchema.to_polars_schema() + assert collection.first.collect_schema() == pl.Schema(MyFirstSchema) assert collection.second is not None assert collection.second.collect().height == 0 - assert collection.second.collect_schema() == MySecondSchema.to_polars_schema() + assert collection.second.collect_schema() == pl.Schema(MySecondSchema) diff --git a/tests/columns/test_polars_schema.py b/tests/columns/test_polars_schema.py index bccf43f0..d63981a5 100644 --- a/tests/columns/test_polars_schema.py +++ b/tests/columns/test_polars_schema.py @@ -9,5 +9,5 @@ def test_polars_schema() -> None: schema = create_schema("test", {"a": dy.Int32(nullable=False), "b": dy.Float32()}) - pl_schema = schema.to_polars_schema() + pl_schema = pl.Schema(schema) assert pl_schema == {"a": pl.Int32, "b": pl.Float32} diff --git a/tests/schema/test_cast.py b/tests/schema/test_cast.py index 3831fccb..e9e0ced8 100644 --- a/tests/schema/test_cast.py +++ b/tests/schema/test_cast.py @@ -28,7 +28,7 @@ def test_cast_valid( df = df_type(data) out = MySchema.cast(df) assert isinstance(out, df_type) - assert out.lazy().collect_schema() == MySchema.to_polars_schema() + assert out.lazy().collect_schema() == pl.Schema(MySchema) def test_cast_invalid_schema_eager() -> None: diff --git a/tests/test_pydantic.py b/tests/test_pydantic.py index ff85d50b..1fa2ce47 100644 --- a/tests/test_pydantic.py +++ b/tests/test_pydantic.py @@ -42,7 +42,7 @@ def df() -> pl.DataFrame: "y": [4, 5, 6], "comment": ["a", "b", "c"], }, - schema=Schema.to_polars_schema(), + schema=pl.Schema(Schema), ) @@ -54,7 +54,7 @@ def invalid_df() -> pl.DataFrame: "y": [4], "comment": ["a"], }, - schema=Schema.to_polars_schema(), + schema=pl.Schema(Schema), )