diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..de288e1 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.formatting.provider": "black" +} \ No newline at end of file diff --git a/env.yml b/env.yml new file mode 100644 index 0000000..6f2e2ca --- /dev/null +++ b/env.yml @@ -0,0 +1,7 @@ +name: typed_polars +channels: + - conda-forge +dependencies: + - python >=3.7 + - polars >=0.15.0 + - pytest >= 7.2.0 \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_simple_df.py b/tests/test_simple_df.py new file mode 100644 index 0000000..39e8877 --- /dev/null +++ b/tests/test_simple_df.py @@ -0,0 +1,19 @@ +import polars as pl +import typed_polars as pt + + +def test_only_types(): + df = pl.DataFrame({"name": ["Hampus", "Dennis", "Noah"], "age": [28, 26, 20]}) + + schema = pt.DataFrameSchema( + {"name": pt.Column(pl.Utf8), "age": pt.Column(pl.Int64)} + ) + + result = schema.validate(df) + + assert result.passed + + df = pl.DataFrame({"age": ["Hampus", "Dennis", "Noah"], "name": [28, 26, 20]}) + result = schema.validate(df) + + assert not result.passed diff --git a/typed_polars/__init__.py b/typed_polars/__init__.py new file mode 100644 index 0000000..f5c4758 --- /dev/null +++ b/typed_polars/__init__.py @@ -0,0 +1,4 @@ +from typed_polars.schema import DataFrameSchema +from typed_polars.column import Column, CastColumn + +import typed_polars.checks as Checks diff --git a/typed_polars/checks/__init__.py b/typed_polars/checks/__init__.py new file mode 100644 index 0000000..f6b80b9 --- /dev/null +++ b/typed_polars/checks/__init__.py @@ -0,0 +1 @@ +# from typed_polars.checks.numerical import le, lt, eq, gt, ge diff --git a/typed_polars/checks/check.py b/typed_polars/checks/check.py new file mode 100644 index 0000000..5064286 --- /dev/null +++ b/typed_polars/checks/check.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass + +import polars as pl + + +# https://pandera.readthedocs.io/en/stable/ +# https://aeturrell.github.io/coding-for-economists/data-advanced.html +# https://pandera.readthedocs.io/en/stable/data_format_conversion.html + + +@dataclass +class ValidationResult: + msg: str + passed: bool + + @staticmethod + def combine_all(results: list[ValidationResult]) -> ValidationResult: + return ValidationResult( + msg="\n".join([m.msg for m in results if m is not None]), + passed=all([m.passed for m in results if m is not None]), + ) + + +class Check(ABC): + @abstractmethod + def validate_column(self, series: pl.Series) -> ValidationResult: + pass + + def validate_frame_column(self, df: pl.DataFrame, column: str) -> ValidationResult: + return self.validate_column(df[column]) diff --git a/typed_polars/checks/numerical.py b/typed_polars/checks/numerical.py new file mode 100644 index 0000000..fe93978 --- /dev/null +++ b/typed_polars/checks/numerical.py @@ -0,0 +1,24 @@ +from datetime import date, datetime +from typing import Callable + +import polars as pl + +def ge(than: float | int | date | datetime) -> Callable[[str], pl.Expr]: + return lambda x: pl.col(x) > than + + + +# @dataclass +# class MustBeGreater(ColumnValidator): +# than: float | int | date | datetime +# inclusive: bool +# +# def validate_frame_column(self, df: pl.DataFrame, column: str): +# col = pl.col(column) +# +# series = df[[column]].filter(col >= self.than if self.inclusive else col > self.than) +# +# return ValidationResult( +# f"Number elements outside bounds {self}: {len(df) - len(series)}", +# len(series) == len(df) +# ) \ No newline at end of file diff --git a/typed_polars/column.py b/typed_polars/column.py new file mode 100644 index 0000000..85a71a3 --- /dev/null +++ b/typed_polars/column.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import polars as pl + +from typed_polars import utils +from typed_polars.checks.check import Check, ValidationResult + + +@dataclass +class Column: + type_: pl.DataType + nullable: bool = False + + checks: list[Check] | Check | None = None + + def check(self, df: pl.DataFrame, column: str) -> ValidationResult: + checks = utils.as_list(self.checks) + results = [] + + # TODO: type check + + for check in checks: + results.append(check.validate_frame_column(df, column)) + + pass + + +@dataclass +class CastColumn(Column): + type_: pl.DataType + nullable: bool = False + + allow_drop_null_on_cast: bool = False + + checks: list[Check] | Check | None = None diff --git a/typed_polars/schema.py b/typed_polars/schema.py new file mode 100644 index 0000000..8c74739 --- /dev/null +++ b/typed_polars/schema.py @@ -0,0 +1,21 @@ +import polars as pl +from dataclasses import dataclass + +from typed_polars.checks.check import ValidationResult +from typed_polars.column import Column + + +@dataclass +class DataFrameSchema: + schema: dict[str, Column] + + def validate(self, df: pl.DataFrame) -> ValidationResult: + results = [] + for column_name, column in self.schema.items(): + column: Column + result = column.check(df, column_name) + results.append(result) + + result: ValidationResult = ValidationResult.combine_all(results) + + return result diff --git a/typed_polars/utils.py b/typed_polars/utils.py new file mode 100644 index 0000000..0b0c095 --- /dev/null +++ b/typed_polars/utils.py @@ -0,0 +1,13 @@ +from typing import TypeVar + +T = TypeVar("T") + + +def as_list(data: list[T] | T | None) -> list[T]: + match data: + case list(): + return data + case None: + return [] + case _: + return [data]