Fix to_yaml serialization dropping global checks

docs/source/schema_inference.rst

@@ -214,6 +214,7 @@ is a convenience method for this functionality.
         coerce: false
         required: true
         regex: false
+    checks: null
     index:
     - pandas_dtype: int64
       nullable: false

pandera/checks.py

@@ -1,11 +1,24 @@
 """Data validation checks."""
+# pylint: disable=fixme
 
 import inspect
 import operator
 import re
-from collections import namedtuple
+from collections import ChainMap, namedtuple
 from functools import partial, wraps
-from typing import Any, Callable, Dict, Iterable, List, Optional, Union
+from itertools import chain
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    no_type_check,
+)
 
 import pandas as pd
 
@@ -397,9 +410,11 @@ def __call__(
         )
 
     def __eq__(self, other):
+        if not isinstance(other, type(self)):
+            return NotImplemented
+
         are_check_fn_objects_equal = (
-            self.__dict__["_check_fn"].__code__.co_code
-            == other.__dict__["_check_fn"].__code__.co_code
+            self._get_check_fn_code() == other._get_check_fn_code()
         )
 
         try:
@@ -427,8 +442,18 @@ def __eq__(self, other):
             and are_all_other_check_attributes_equal
         )
 
+    def _get_check_fn_code(self):
+        check_fn = self.__dict__["_check_fn"]
+        try:
+            code = check_fn.__code__.co_code
+        except AttributeError:
+            # try accessing the functools.partial wrapper
+            code = check_fn.func.__code__.co_code
+
+        return code
+
     def __hash__(self):
-        return hash(self.__dict__["_check_fn"].__code__.co_code)
+        return hash(self._get_check_fn_code())
 
     def __repr__(self):
         return (
@@ -438,22 +463,49 @@ def __repr__(self):
         )
 
 
+_T = TypeVar("_T", bound=_CheckBase)
+
+
 class _CheckMeta(type):  # pragma: no cover
     """Check metaclass."""
 
+    # FIXME: this should probably just be moved to _CheckBase
+
+    REGISTERED_CUSTOM_CHECKS: Dict[str, Callable] = {}  # noqa
+
     def __getattr__(cls, name: str) -> Any:
         """Prevent attribute errors for registered checks."""
-        attr = cls.__dict__.get(name)
+        attr = ChainMap(cls.__dict__, cls.REGISTERED_CUSTOM_CHECKS).get(name)
         if attr is None:
-            raise AttributeError(f"'{cls}' object has no attribute '{name}'")
+            raise AttributeError(
+                f"'{cls}' object has no attribute '{name}'. "
+                "Make sure any custom checks have been registered "
+                "using the extensions api."
+            )
         return attr
 
+    def __dir__(cls) -> Iterable[str]:
+        """Allow custom checks to show up as attributes when autocompleting."""
+        return chain(super().__dir__(), cls.REGISTERED_CUSTOM_CHECKS.keys())
+
+    # pylint: disable=line-too-long
+    # mypy has limited metaclass support so this doesn't pass typecheck
+    # see https://mypy.readthedocs.io/en/stable/metaclasses.html#gotchas-and-limitations-of-metaclass-support
+    # pylint: enable=line-too-long
+    @no_type_check
+    def __contains__(cls: Type[_T], item: Union[_T, str]) -> bool:
+        """Allow lookups for registered checks."""
+        if isinstance(item, cls):
+            name = item.name
+            return hasattr(cls, name)
+
+        # assume item is str
+        return hasattr(cls, item)
+
 
 class Check(_CheckBase, metaclass=_CheckMeta):
     """Check a pandas Series or DataFrame for certain properties."""
 
-    REGISTERED_CUSTOM_CHECKS: Dict[str, Callable] = {}  # noqa
-
     @classmethod
     @st.register_check_strategy(st.eq_strategy)
     @register_check_statistics(["value"])

pandera/extensions.py

@@ -161,9 +161,8 @@ def check_method(cls, *args, **kwargs):
         if strategy is not None:
             check_method = st.register_check_strategy(strategy)(check_method)
 
-        setattr(Check, check_fn.__name__, classmethod(check_method))
-        Check.REGISTERED_CUSTOM_CHECKS[check_fn.__name__] = getattr(
-            Check, check_fn.__name__
+        Check.REGISTERED_CUSTOM_CHECKS[check_fn.__name__] = partial(
+            check_method, Check
         )
 
     return register_check_wrapper(check_fn)

pandera/io.py

@@ -25,40 +25,60 @@
 NOT_JSON_SERIALIZABLE = {PandasDtype.DateTime, PandasDtype.Timedelta}
 
 
-def _serialize_check_stats(check_stats, pandas_dtype):
+def _serialize_check_stats(check_stats, pandas_dtype=None):
     """Serialize check statistics into json/yaml-compatible format."""
 
-    def handle_stat_dtype(stat):
+    def handle_stat_dtype(stat, pandas_dtype):
         if pandas_dtype == PandasDtype.DateTime:
             return stat.strftime(DATETIME_FORMAT)
         elif pandas_dtype == PandasDtype.Timedelta:
             # serialize to int in nanoseconds
             return stat.delta
+
         return stat
 
     # for unary checks, return a single value instead of a dictionary
     if len(check_stats) == 1:
-        return handle_stat_dtype(list(check_stats.values())[0])
+        return handle_stat_dtype(list(check_stats.values())[0], pandas_dtype)
 
     # otherwise return a dictionary of keyword args needed to create the Check
     serialized_check_stats = {}
     for arg, stat in check_stats.items():
-        serialized_check_stats[arg] = handle_stat_dtype(stat)
+        serialized_check_stats[arg] = handle_stat_dtype(stat, pandas_dtype)
     return serialized_check_stats
 
 
+def _serialize_dataframe_stats(dataframe_checks):
+    """
+    Serialize global dataframe check statistics into json/yaml-compatible format.
+    """
+    serialized_checks = {}
+
+    for check_name, check_stats in dataframe_checks.items():
+        # The case that `check_name` is not registered is handled in `parse_checks`,
+        # so we know that `check_name` exists.
+
+        # infer dtype of statistics and serialize them
+        serialized_checks[check_name] = _serialize_check_stats(check_stats)
+
+    return serialized_checks
+
+
 def _serialize_component_stats(component_stats):
     """
     Serialize column or index statistics into json/yaml-compatible format.
     """
+    # pylint: disable=import-outside-toplevel
+    from pandera.checks import Check
+
     serialized_checks = None
     if component_stats["checks"] is not None:
         serialized_checks = {}
         for check_name, check_stats in component_stats["checks"].items():
-            if check_stats is None:
+            if check_name not in Check:
                 warnings.warn(
                     f"Check {check_name} cannot be serialized. This check will be "
-                    f"ignored"
+                    "ignored. Did you forget to register it with the extension API?"
                 )
             else:
                 serialized_checks[check_name] = _serialize_check_stats(
@@ -93,7 +113,7 @@ def _serialize_schema(dataframe_schema):
 
     statistics = get_dataframe_schema_statistics(dataframe_schema)
 
-    columns, index = None, None
+    columns, index, checks = None, None, None
     if statistics["columns"] is not None:
         columns = {
             col_name: _serialize_component_stats(column_stats)
@@ -106,18 +126,22 @@ def _serialize_schema(dataframe_schema):
             for index_stats in statistics["index"]
         ]
 
+    if statistics["checks"] is not None:
+        checks = _serialize_dataframe_stats(statistics["checks"])
+
     return {
         "schema_type": "dataframe",
         "version": __version__,
         "columns": columns,
+        "checks": checks,
         "index": index,
         "coerce": dataframe_schema.coerce,
         "strict": dataframe_schema.strict,
     }
 
 
-def _deserialize_check_stats(check, serialized_check_stats, pandas_dtype):
-    def handle_stat_dtype(stat):
+def _deserialize_check_stats(check, serialized_check_stats, pandas_dtype=None):
+    def handle_stat_dtype(stat, pandas_dtype):
         if pandas_dtype == PandasDtype.DateTime:
             return pd.to_datetime(stat, format=DATETIME_FORMAT)
         elif pandas_dtype == PandasDtype.Timedelta:
@@ -130,10 +154,10 @@ def handle_stat_dtype(stat):
         # dictionary mapping Check arg names to values.
         check_stats = {}
         for arg, stat in serialized_check_stats.items():
-            check_stats[arg] = handle_stat_dtype(stat)
+            check_stats[arg] = handle_stat_dtype(stat, pandas_dtype)
         return check(**check_stats)
     # otherwise assume unary check function signature
-    return check(handle_stat_dtype(serialized_check_stats))
+    return check(handle_stat_dtype(serialized_check_stats, pandas_dtype))
 
 
 def _deserialize_component_stats(serialized_component_stats):
@@ -173,9 +197,9 @@ def _deserialize_component_stats(serialized_component_stats):
 
 def _deserialize_schema(serialized_schema):
     # pylint: disable=import-outside-toplevel
-    from pandera import Column, DataFrameSchema, Index, MultiIndex
+    from pandera import Check, Column, DataFrameSchema, Index, MultiIndex
 
-    columns, index = None, None
+    columns, index, checks = None, None, None
     if serialized_schema["columns"] is not None:
         columns = {
             col_name: Column(**_deserialize_component_stats(column_stats))
@@ -188,6 +212,13 @@ def _deserialize_schema(serialized_schema):
             for index_component in serialized_schema["index"]
         ]
 
+    if serialized_schema["checks"] is not None:
+        # handles unregistered checks by raising AttributeErrors from getattr
+        checks = [
+            _deserialize_check_stats(getattr(Check, check_name), check_stats)
+            for check_name, check_stats in serialized_schema["checks"].items()
+        ]
+
     if index is None:
         pass
     elif len(index) == 1:
@@ -199,6 +230,7 @@ def _deserialize_schema(serialized_schema):
 
     return DataFrameSchema(
         columns=columns,
+        checks=checks,
         index=index,
         coerce=serialized_schema["coerce"],
         strict=serialized_schema["strict"],

pandera/model.py

@@ -1,5 +1,6 @@
 """Class-based api"""
 import inspect
+import os
 import re
 import sys
 import typing
@@ -170,6 +171,13 @@ def to_schema(cls) -> DataFrameSchema:
             MODEL_CACHE[cls] = cls.__schema__
         return cls.__schema__
 
+    @classmethod
+    def to_yaml(cls, stream: Optional[os.PathLike] = None):
+        """
+        Convert `Schema` to yaml using `io.to_yaml`.
+        """
+        return cls.to_schema().to_yaml(stream)
+
     @classmethod
     @pd.util.Substitution(validate_doc=DataFrameSchema.validate.__doc__)
     def validate(

pandera/schema_statistics.py

@@ -115,6 +115,7 @@ def get_dataframe_schema_statistics(dataframe_schema):
             }
             for col_name, column in dataframe_schema.columns.items()
         },
+        "checks": parse_checks(dataframe_schema.checks),
         "index": (
             None
             if dataframe_schema.index is None
@@ -158,7 +159,17 @@ def parse_checks(checks) -> Union[Dict[str, Any], None]:
     check_statistics = {}
     _check_memo = {}
     for check in checks:
-        check_statistics[check.name] = check.statistics
+        if check not in Check:
+            warnings.warn(
+                "Only registered checks may be serialized to statistics. "
+                "Did you forget to register it with the extension API? "
+                f"Check `{check.name}` will be skipped."
+            )
+            continue
+
+        check_statistics[check.name] = (
+            {} if check.statistics is None else check.statistics
+        )
         _check_memo[check.name] = check
 
     # raise ValueError on incompatible checks

pandera/schemas.py

@@ -3,6 +3,7 @@
 
 import copy
 import itertools
+import os
 import warnings
 from functools import wraps
 from pathlib import Path
@@ -1186,17 +1187,16 @@ def from_yaml(cls, yaml_schema) -> "DataFrameSchema":
 
         return pandera.io.from_yaml(yaml_schema)
 
-    def to_yaml(self, fp: Union[str, Path] = None):
+    def to_yaml(self, stream: Optional[os.PathLike] = None):
         """Write DataFrameSchema to yaml file.
 
-        :param dataframe_schema: schema to write to file or dump to string.
         :param stream: file stream to write to. If None, dumps to string.
         :returns: yaml string if stream is None, otherwise returns None.
         """
         # pylint: disable=import-outside-toplevel,cyclic-import
         import pandera.io
 
-        return pandera.io.to_yaml(self, fp)
+        return pandera.io.to_yaml(self, stream=stream)
 
     def set_index(
         self, keys: List[str], drop: bool = True, append: bool = False

tests/core/checks_fixtures.py

@@ -0,0 +1,33 @@
+"""Pytest fixtures for testing custom checks."""
+import unittest.mock as mock
+
+import pandas as pd
+import pytest
+
+import pandera as pa
+import pandera.extensions as pa_ext
+
+__all__ = "custom_check_teardown", "extra_registered_checks"
+
+
+@pytest.fixture(scope="function")
+def custom_check_teardown():
+    """Remove all custom checks after execution of each pytest function."""
+    yield
+    for check_name in list(pa.Check.REGISTERED_CUSTOM_CHECKS):
+        del pa.Check.REGISTERED_CUSTOM_CHECKS[check_name]
+
+
+@pytest.fixture(scope="function")
+def extra_registered_checks():
+    """temporarily registers custom checks onto the Check class"""
+    # pylint: disable=unused-variable
+    with mock.patch(
+        "pandera.Check.REGISTERED_CUSTOM_CHECKS", new_callable=dict
+    ):
+        # register custom checks here
+        @pa_ext.register_check_method()
+        def no_param_check(_: pd.DataFrame) -> bool:
+            return True
+
+        yield

tests/core/conftest.py

@@ -0,0 +1,4 @@
+"""Registers fixtures for core"""
+
+# pylint: disable=unused-import
+from .checks_fixtures import custom_check_teardown, extra_registered_checks