Serializing custom checks

[cosmicBboy]

This post answers a question by @antonl, originally posted here: #383 (comment)

With the new extensions API, you can register custom checks into the Check namespace, which gives you access as a method Check.my_check(...) or as a keyword argument in the class-based API column: Series[int] = Field(my_check=...)

Here's an example script:

import pandera as pa
import pandera.extensions as extensions
import pandas as pd


@extensions.register_check_method(statistics=["min_value", "max_value"])
def is_between(pandas_obj, *, min_value, max_value):
    return (min_value <= pandas_obj) & (pandas_obj <= max_value)


schema = pa.DataFrameSchema(
    {"col": pa.Column(pa.Int, pa.Check.is_between(min_value=1, max_value=10))}
)

data = pd.DataFrame({"col": [1, 5, 10]})

yaml_schema = schema.to_yaml()
print("yaml-serialized schema with custom check")
print("----------------------------------------")
print(yaml_schema)
schema_from_yaml = pa.io.from_yaml(yaml_schema)

print("\nschema object from yaml")
print("------------------------")
print(schema_from_yaml)

print("\ncustom check is persisted")
print("--------------------------")
print(schema_from_yaml.columns["col"].checks)

print("\n validated dataframe")
print("--------------------------")
print(schema_from_yaml(pd.DataFrame({"col": [1, 5, 10]})))

Output:

yaml-serialized schema with custom check
----------------------------------------
schema_type: dataframe
version: 0.6.2
columns:
  col:
    pandas_dtype: int
    nullable: false
    checks:
      is_between:
        min_value: 1
        max_value: 10
    allow_duplicates: true
    coerce: false
    required: true
    regex: false
index: null
coerce: false
strict: false


schema object from yaml
------------------------
<Schema DataFrameSchema(
    columns={
        'col': <Schema Column(name=col, type=int)>
    },
    checks=[],
    coerce=False,
    pandas_dtype=None,
    index=None,
    strict=False
    name=None,
    ordered=False
)>

custom check is persisted
--------------------------
[<Check is_between>]

 validated dataframe
--------------------------
   col
0    1
1    5
2   10

[antonl]

I agree, this works, but I was hoping to do global checks. I know this is a relatively new part of pandera, but I believe the current behavior of to_yaml has a bug. Here's a test case. edit: posted as bug #419 .

import pandera
import pandera.extensions as extensions
import pandera.typing as pat


class GlobalSerdes(pandera.SchemaModel):
    """initial schema that includes a global check"""
    a: pat.Series[pat.Int64]
    b: pat.Series[pat.Int64]

    @pandera.dataframe_check()
    def _a_gt_b(self, df):
        """ensure that every `a` is greater than every `b`."""
        return (df["a"] > df["b"]).all()


class PostSerdes(pandera.SchemaModel):
    """after serdes, it looks like this"""
    a: pat.Series[pat.Int64]
    b: pat.Series[pat.Int64]


def test_schema_model_serdes():
    """demonstrate that to_yaml strips custom global checks"""
    initial = GlobalSerdes.to_schema().to_yaml()
    final = PostSerdes.to_schema().to_yaml()

    assert initial != final, "global check was stripped"



def test_data_frame_schema():
    """demonstrate that to_yaml fails to serialize registered global checks"""

    @extensions.register_check_method(statistics=["a_gt_b"])
    def is_a_gt_b(pandas_obj, *, a_gt_b: bool):
        return (pandas_obj["a"] > pandas_obj["b"]).all() == a_gt_b

    initial = PostSerdes.to_schema().to_yaml()

    _ = PostSerdes.to_schema()
    _.checks.append(pandera.Check.is_a_gt_b(a_gt_b=True))

    final = _.to_yaml()

    assert initial != final, "registered global custom check failed"


if __name__ == "__main__":
    test_schema_model_serdes()
    test_data_frame_schema()