First Code (ALPHA)

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.formatting.provider": "black"
+}

env.yml

@@ -0,0 +1,7 @@
+name: typed_polars
+channels:
+  - conda-forge
+dependencies:
+  - python >=3.7
+  - polars >=0.15.0
+  - pytest >= 7.2.0

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

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

typed_polars/checks/__init__.py

@@ -0,0 +1 @@
+# from typed_polars.checks.numerical import le, lt, eq, gt, ge

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])

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)
+#         )

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

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

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]