feat: add schema generation (lazy and eager) (#11)

* initial schema generation

* return repr object for easier io

* rename optional dependency to "pyperclip"

* support lazy schema generation

* make __repr__ return the schema string

* comments

* tests

Author: CangyuanLi

pyproject.toml

@@ -11,6 +11,10 @@ classifiers = [
 ]
 dependencies = ["narwhals>=1.0.0"]
 
+[project.optional-dependencies]
+pyperclip = ["pyperclip>=1.0.0"]
+
+
 [build-system]
 requires = ["setuptools>=42", "wheel"]
 build-backend = "setuptools.build_meta"

src/checkedframe/__init__.py

@@ -31,3 +31,4 @@
     UInt128,
     Unknown,
 )
+from ._schema_generation import generate_schema_repr

src/checkedframe/_dtypes.py

@@ -2,7 +2,7 @@
 
 from abc import abstractmethod
 from collections.abc import Mapping
-from typing import TYPE_CHECKING, Literal, Optional, Union
+from typing import TYPE_CHECKING, Literal, Optional, TypedDict, Union
 
 import narwhals.stable.v1 as nw
 from narwhals.stable.v1.dtypes import DType as NarwhalsDType
@@ -30,6 +30,14 @@ class _BoundedDType(_DType):
     _max: int | float
 
 
+class _ColumnKwargs(TypedDict):
+    name: Optional[str]
+    nullable: bool
+    required: bool
+    cast: bool
+    checks: Optional[list[Check]]
+
+
 class _Column:
     """Represents a column in a DataFrame.
 
@@ -945,11 +953,15 @@ def __to_narwhals(self):
         return nw.Datetime(time_unit=self.time_unit, time_zone=self.time_zone)
 
     @staticmethod
-    def from_narwhals(nw_dtype: nw.Datetime) -> Datetime:
+    def from_narwhals(nw_dtype: nw.Datetime, **column_kwargs) -> Datetime:
         if hasattr(nw_dtype, "time_unit"):
-            return Datetime(time_unit=nw_dtype.time_unit, time_zone=nw_dtype.time_zone)
+            return Datetime(
+                time_unit=nw_dtype.time_unit,
+                time_zone=nw_dtype.time_zone,
+                **column_kwargs,
+            )
 
-        return Datetime()
+        return Datetime(**column_kwargs)
 
     @staticmethod
     def _safe_cast(s: nw.Series, to_dtype: _DType) -> nw.Series:
@@ -987,11 +999,11 @@ def __to_narwhals(self):
         return nw.Duration(time_unit=self.time_unit)
 
     @staticmethod
-    def from_narwhals(nw_dtype: nw.Duration) -> Duration:
+    def from_narwhals(nw_dtype: nw.Duration, **column_kwargs) -> Duration:
         if hasattr(nw_dtype, "time_unit"):
-            return Duration(nw_dtype.time_unit)
+            return Duration(nw_dtype.time_unit, **column_kwargs)
 
-        return Duration()
+        return Duration(**column_kwargs)
 
     @staticmethod
     def _safe_cast(s: nw.Series, to_dtype: _DType) -> nw.Series:
@@ -1113,8 +1125,10 @@ def to_narwhals(self):  # type: ignore
         return nw.Array(self.inner.to_narwhals(), self.shape)
 
     @staticmethod
-    def from_narwhals(nw_dtype: nw.Array) -> Array:
-        return Array(_nw_type_to_cf_type(nw_dtype.inner), shape=nw_dtype.shape)
+    def from_narwhals(nw_dtype: nw.Array, **column_kwargs) -> Array:
+        return Array(
+            _nw_type_to_cf_type(nw_dtype.inner), shape=nw_dtype.shape, **column_kwargs
+        )
 
     @staticmethod
     def _safe_cast(s: nw.Series, to_dtype: _DType) -> nw.Series:
@@ -1148,8 +1162,8 @@ def to_narwhals(self):  # type: ignore[override]
         return nw.List(self.inner.to_narwhals())
 
     @staticmethod
-    def from_narwhals(nw_dtype: nw.List) -> List:
-        return List(_nw_type_to_cf_type(nw_dtype.inner))
+    def from_narwhals(nw_dtype: nw.List, **column_kwargs) -> List:
+        return List(_nw_type_to_cf_type(nw_dtype.inner), **column_kwargs)
 
     @staticmethod
     def _safe_cast(s: nw.Series, to_dtype: _DType) -> nw.Series:
@@ -1196,56 +1210,58 @@ def to_narwhals(self) -> nw.Struct:  # type: ignore
         return nw.Struct(dct)
 
     @staticmethod
-    def from_narwhals(nw_dtype: nw.Struct) -> Struct:
+    def from_narwhals(nw_dtype: nw.Struct, **column_kwargs) -> Struct:
         dct = {}
         for field in nw_dtype.fields:
             dct[field.name] = _nw_type_to_cf_type(field.dtype)
 
-        return Struct(dct)
+        return Struct(dct, **column_kwargs)
 
     @staticmethod
     def _safe_cast(s: nw.Series, to_dtype: _DType) -> nw.Series:
         return _checked_cast(s, to_dtype)
 
 
-_NARWHALS_DTYPE_TO_CHECKEDFRAME_DTYPE_MAPPER: dict[type[NarwhalsDType], _DType] = {
-    nw.Binary: Binary(),
-    nw.Boolean: Boolean(),
-    nw.Categorical: Categorical(),
-    nw.Date: Date(),
-    nw.Datetime: Datetime(),
-    nw.Decimal: Decimal(),
-    nw.Enum: Enum(),
-    nw.Float32: Float32(),
-    nw.Float64: Float64(),
-    nw.Int8: Int8(),
-    nw.Int16: Int16(),
-    nw.Int32: Int32(),
-    nw.Int64: Int64(),
-    nw.Int128: Int128(),
-    nw.Object: Object(),
-    nw.String: String(),
-    nw.UInt8: UInt8(),
-    nw.UInt16: UInt16(),
-    nw.UInt32: UInt32(),
-    nw.UInt64: UInt64(),
-    nw.UInt128: UInt128(),
-    nw.Unknown: Unknown(),
+_NARWHALS_DTYPE_TO_CHECKEDFRAME_DTYPE_MAPPER: dict[
+    type[NarwhalsDType], type[_DType]
+] = {
+    nw.Binary: Binary,
+    nw.Boolean: Boolean,
+    nw.Categorical: Categorical,
+    nw.Date: Date,
+    nw.Datetime: Datetime,
+    nw.Decimal: Decimal,
+    nw.Enum: Enum,
+    nw.Float32: Float32,
+    nw.Float64: Float64,
+    nw.Int8: Int8,
+    nw.Int16: Int16,
+    nw.Int32: Int32,
+    nw.Int64: Int64,
+    nw.Int128: Int128,
+    nw.Object: Object,
+    nw.String: String,
+    nw.UInt8: UInt8,
+    nw.UInt16: UInt16,
+    nw.UInt32: UInt32,
+    nw.UInt64: UInt64,
+    nw.UInt128: UInt128,
+    nw.Unknown: Unknown,
 }
 
 
 def _nw_type_to_cf_type(
-    nw_dtype: Union[NarwhalsDType, type[NarwhalsDType]],
+    nw_dtype: Union[NarwhalsDType, type[NarwhalsDType]], **column_kwargs
 ) -> _DType:
     if isinstance(nw_dtype, nw.Array):
-        return Array.from_narwhals(nw_dtype)
+        return Array.from_narwhals(nw_dtype, **column_kwargs)
     elif isinstance(nw_dtype, nw.List):
-        return List.from_narwhals(nw_dtype)
+        return List.from_narwhals(nw_dtype, **column_kwargs)
     elif isinstance(nw_dtype, nw.Struct):
-        return Struct.from_narwhals(nw_dtype)
+        return Struct.from_narwhals(nw_dtype, **column_kwargs)
     elif isinstance(nw_dtype, nw.Datetime):
-        return Datetime.from_narwhals(nw_dtype)
+        return Datetime.from_narwhals(nw_dtype, **column_kwargs)
     elif isinstance(nw_dtype, nw.Duration):
-        return Duration.from_narwhals(nw_dtype)
+        return Duration.from_narwhals(nw_dtype, **column_kwargs)
 
-    return _NARWHALS_DTYPE_TO_CHECKEDFRAME_DTYPE_MAPPER[nw_dtype]  # type: ignore
+    return _NARWHALS_DTYPE_TO_CHECKEDFRAME_DTYPE_MAPPER[nw_dtype](**column_kwargs)  # type: ignore

src/checkedframe/_schema_generation.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import keyword
+from pathlib import Path
+from typing import Optional
+
+import narwhals.stable.v1 as nw
+import narwhals.stable.v1.selectors as nws
+import narwhals.stable.v1.typing as nwt
+
+from ._dtypes import _nw_type_to_cf_type
+
+INF = float("inf")
+NEG_INF = float("-inf")
+
+
+class SchemaRepr:
+    def __init__(self, schema_repr: str):
+        self.schema_repr = schema_repr
+
+    def __repr__(self) -> str:
+        # This is useful for usage in notebooks, where the user may want to just copy
+        # the output from the notebook cell.
+        return self.schema_repr
+
+    def write_clipboard(self):
+        import pyperclip  # type: ignore
+
+        pyperclip.copy(self.schema_repr)
+
+    def write_text(self, file: str | Path):
+        Path(file).write_text(self.schema_repr)
+
+
+def generate_schema_repr(
+    df: nwt.IntoFrame,
+    lazy: bool = False,
+    class_name: str = "MySchema",
+    header: Optional[str] = "import checkedframe as cf",
+    import_alias: str = "cf.",
+) -> SchemaRepr:
+    nw_df = nw.from_native(df)
+
+    if isinstance(nw_df, nw.LazyFrame):
+        lazy = True
+
+    if not lazy:
+        null_df = nw_df.select(nws.all().is_null().any())
+
+        float_selector = nws.by_dtype(nw.Float32, nw.Float64)
+        nan_df = nw_df.select(float_selector.is_nan().any())
+        inf_df = nw_df.select(float_selector.is_in((INF, NEG_INF)).any())
+
+        float_cols = set(nan_df.columns)
+
+    # Build string representation of schema
+    columns = []
+    i = 0
+    for col, nw_dtype in nw_df.collect_schema().items():
+        column_kwargs: dict[str, bool | str] = {}
+        if not col.isidentifier() or keyword.iskeyword(col):
+            column_kwargs["name"] = f'"{col}"'
+
+            sanitized_col = f"column_{i}"
+            i += 1
+        else:
+            sanitized_col = col
+
+        if not lazy:
+            if null_df[col].item():
+                column_kwargs["nullable"] = True
+
+            if col in float_cols:
+                if nan_df[col].item():
+                    column_kwargs["allow_nan"] = True
+
+                if inf_df[col].item():
+                    column_kwargs["allow_inf"] = True
+
+        cf_dtype = _nw_type_to_cf_type(nw_dtype, **column_kwargs)
+
+        kwargs_to_show = []
+        for k, v in column_kwargs.items():
+            kwargs_to_show.append(f"{k}={v}")
+
+        display_kwargs = ", ".join(kwargs_to_show)
+        display_dtype = str(cf_dtype).replace("(", f"({import_alias}")
+
+        columns.append(
+            f"    {sanitized_col} = {import_alias}{display_dtype}({display_kwargs})".replace(
+                ")(", ", " if len(kwargs_to_show) > 0 else ""
+            )
+        )
+
+    if header is not None:
+        header = f"{header}\n\n"
+    else:
+        header = ""
+
+    col_repr = "\n".join(columns)
+
+    schema_repr = f"{header}class {class_name}({import_alias}Schema):\n{col_repr}"
+
+    return SchemaRepr(schema_repr)

tests/test_schema_parsing.py

@@ -1,3 +1,6 @@
+import tempfile
+
+import polars as pl
 import pytest
 
 import checkedframe as cf
@@ -91,3 +94,58 @@ def a_check() -> bool:
     assert name in schema_dict
 
     assert len(schema_dict[name].checks) == 1
+
+
+def test_schema_generation():
+    correct_repr = (
+        "import checkedframe as cf\n"
+        "\n"
+        "class AASchema(cf.Schema):\n"
+        '    column_0 = cf.Float64(name="reason code", nullable=True, allow_nan=True)\n'
+        "    y = cf.List(cf.Int64, nullable=True)\n"
+        "    z = cf.List(cf.List(cf.Int64), nullable=True)"
+    )
+
+    schema_repr = cf.generate_schema_repr(
+        pl.DataFrame(
+            {
+                "reason code": [1.0, float("nan"), None],
+                "y": [[1], [2], None],
+                "z": [[[1]], None, [[3]]],
+            }
+        ),
+        class_name="AASchema",
+    )
+
+    assert schema_repr.schema_repr == correct_repr
+
+
+def test_schema_generation_lazy():
+    correct_repr = (
+        "import checkedframe as cf\n"
+        "\n"
+        "class AASchema(cf.Schema):\n"
+        '    column_0 = cf.Float64(name="reason code")\n'
+        "    y = cf.List(cf.Int64)\n"
+        "    z = cf.List(cf.List(cf.Int64))"
+    )
+
+    df = pl.DataFrame(
+        {
+            "reason code": [1.0, float("nan"), None],
+            "y": [[1], [2], None],
+            "z": [[[1]], None, [[3]]],
+        }
+    )
+
+    schema_repr = cf.generate_schema_repr(
+        df,
+        lazy=True,
+        class_name="AASchema",
+    )
+
+    assert schema_repr.schema_repr == correct_repr
+
+    schema_repr = cf.generate_schema_repr(df.lazy(), class_name="AASchema")
+
+    assert schema_repr.schema_repr == correct_repr