Repository: unionai-oss/pandera

Issue:
Make ``serialize_schema`` public or add ``to_json``/``from_json``
**Is your feature request related to a problem? Please describe.**
We would like to store the pandera schema as a JSON file for the compatibility with the downstream tools.
However, the current API allows only for YAML format.

**Describe the solution you'd like**
The easiest and probably the most general solution would be to make ``pandera.io._serialize_schema`` public by merely renaming it.

Here is the relevant line: https://github.com/unionai-oss/pandera/blob/31b064352d2d6173147bfd9389066028eed385d6/pandera/io.py#L126

This would allow the user to serialize it to JSON, but also to any other JSON-like format.

Alternatively, if you do not want to expose `_serialize_schema`, you could add ``to_json``/``from_json`` methods, analogous to YAML.

**Describe alternatives you've considered**
There is no nice alternative. We could parse YAML file, and then convert it, but that would force us to add YAML as a dependency in our code.

Thank you very much for your work on this library! Please let us know if we can help in any way.


Discussion:
hi @KiaXdice, this use-case makes sense!

I think a `to_json` / `from_json` method makes a lot of sense, would you be open to making a PR for that? Should be fairly straight-forward, as the output of `_serialize_schema` is already a json-serializable (I think?) python dict 

> but that would force us to add YAML as a dependency in our code.

so `pip install 'pandera[io]'` will actually install the `pyyaml` package... is this an issue for you, or do you just not want it as an explicit dependency in your code?
@cosmicBboy thanks, we'll try to make a PR tomorrow (CET).

> so pip install 'pandera[io]' will actually install the pyyaml package... is this an issue for you, or do you just not want it as an explicit dependency in your code?

Installing the dependency implicitly/by transitivity is OK. We just prefer to have no direct YAML dependency in our code.


Resolving patch:
diff --git a/pandera/io.py b/pandera/io.py
--- a/pandera/io.py
+++ b/pandera/io.py
@@ -1,5 +1,6 @@
 """Module for reading and writing schema objects."""
 
+import json
 import warnings
 from collections.abc import Mapping
 from functools import partial
@@ -123,7 +124,7 @@ def _serialize_component_stats(component_stats):
     }
 
 
-def _serialize_schema(dataframe_schema):
+def serialize_schema(dataframe_schema):
     """Serialize dataframe schema into into json/yaml-compatible format."""
     from pandera import __version__  # pylint: disable=import-outside-toplevel
 
@@ -217,7 +218,14 @@ def _deserialize_component_stats(serialized_component_stats):
     }
 
 
-def _deserialize_schema(serialized_schema):
+def deserialize_schema(serialized_schema):
+    """
+    De-serialize the schema from a JSON-able mapping.
+
+    :param serialized_schema: a mapping representing the schema
+    :returns:
+        the schema de-serialized into :class:`~pandera.schemas.DataFrameSchema`
+    """
     # pylint: disable=import-outside-toplevel
     from pandera import Index, MultiIndex
 
@@ -282,7 +290,7 @@ def from_yaml(yaml_schema):
             serialized_schema = yaml.safe_load(f)
     except (TypeError, OSError):
         serialized_schema = yaml.safe_load(yaml_schema)
-    return _deserialize_schema(serialized_schema)
+    return deserialize_schema(serialized_schema)
 
 
 def to_yaml(dataframe_schema, stream=None):
@@ -292,7 +300,7 @@ def to_yaml(dataframe_schema, stream=None):
     :param stream: file stream to write to. If None, dumps to string.
     :returns: yaml string if stream is None, otherwise returns None.
     """
-    statistics = _serialize_schema(dataframe_schema)
+    statistics = serialize_schema(dataframe_schema)
 
     def _write_yaml(obj, stream):
         return yaml.safe_dump(obj, stream=stream, sort_keys=False)
@@ -304,6 +312,57 @@ def _write_yaml(obj, stream):
         return _write_yaml(statistics, stream)
 
 
+def from_json(source):
+    """
+    Create :class:`~pandera.schemas.DataFrameSchema` from json file.
+
+    :param source:
+        Depending on the type, source is assumed to be:
+
+        1) str or Path to a file containing json schema (if the file exists),
+        2) str as a JSON-encoded schema, or
+        3) stream that we can read from containing the schema as JSON-encoded
+           string.
+
+    :returns: dataframe schema.
+    """
+    if isinstance(source, str):
+        try:
+            serialized_schema = json.loads(source)
+        except json.decoder.JSONDecodeError:
+            with Path(source).open(encoding="utf-8") as f:
+                serialized_schema = json.load(fp=f)
+    elif isinstance(source, Path):
+        with source.open(encoding="utf-8") as f:
+            serialized_schema = json.load(fp=f)
+    else:
+        serialized_schema = json.load(fp=source)
+
+    return deserialize_schema(serialized_schema)
+
+
+def to_json(dataframe_schema, target=None, **kwargs):
+    """
+    Write :class:`~pandera.schemas.DataFrameSchema` to json file.
+
+    :param dataframe_schema: schema to write to file or dump to string.
+    :param target: file path or stream to write to. If None, returns a
+        dump to string.
+    :param kwargs: keyword arguments to pass into :func:`json.dump`
+    :returns: json string if stream is None, otherwise returns None.
+    """
+    serialized_schema = serialize_schema(dataframe_schema)
+
+    if target is None:
+        return json.dumps(serialized_schema, sort_keys=False, **kwargs)
+
+    if isinstance(target, (str, Path)):
+        with Path(target).open("w", encoding="utf-8") as f:
+            json.dump(serialized_schema, fp=f, sort_keys=False, **kwargs)
+    else:
+        json.dump(serialized_schema, fp=target, sort_keys=False, **kwargs)
+
+
 SCRIPT_TEMPLATE = """
 from pandera import (
     DataFrameSchema, Column, Check, Index, MultiIndex
@@ -705,4 +764,4 @@ def from_frictionless_schema(
             None if len(schema.primary_key) == 1 else list(schema.primary_key)
         ),
     }
-    return _deserialize_schema(assembled_schema)
+    return deserialize_schema(assembled_schema)
diff --git a/pandera/schemas.py b/pandera/schemas.py
--- a/pandera/schemas.py
+++ b/pandera/schemas.py
@@ -1377,13 +1377,52 @@ def to_yaml(self, stream: os.PathLike) -> None:  # pragma: no cover
     def to_yaml(self, stream: Optional[os.PathLike] = None) -> Optional[str]:
         """Write DataFrameSchema to yaml file.
 
-        :param stream: file stream to write to. If None, dumps to string.
+        :param stream: file path or 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, stream=stream)
+        return pandera.io.to_yaml(self, stream)
+
+    @classmethod
+    def from_json(cls, source) -> "DataFrameSchema":
+        """Create DataFrameSchema from json file.
+
+        :param source: str, Path to json schema, or serialized yaml
+            string.
+        :returns: dataframe schema.
+        """
+        # pylint: disable=import-outside-toplevel,cyclic-import
+        import pandera.io
+
+        return pandera.io.from_json(source)
+
+    @overload
+    def to_json(
+        self, target: None = None, **kwargs
+    ) -> str:  # pragma: no cover
+        ...
+
+    @overload
+    def to_json(
+        self, target: os.PathLike, **kwargs
+    ) -> None:  # pragma: no cover
+        ...
+
+    def to_json(
+        self, target: Optional[os.PathLike] = None, **kwargs
+    ) -> Optional[str]:
+        """Write DataFrameSchema to json file.
+
+        :param target: file target to write to. If None, dumps to string.
+        :returns: json string if target is None, otherwise returns None.
+        """
+        # pylint: disable=import-outside-toplevel,cyclic-import
+        import pandera.io
+
+        return pandera.io.to_json(self, target, **kwargs)
 
     def set_index(
         self, keys: List[str], drop: bool = True, append: bool = False


Test patch:
diff --git a/tests/io/test_io.py b/tests/io/test_io.py
--- a/tests/io/test_io.py
+++ b/tests/io/test_io.py
@@ -2,6 +2,7 @@
 
 import platform
 import tempfile
+from io import StringIO
 from pathlib import Path
 from unittest import mock
 
@@ -701,7 +702,7 @@ def test_io_yaml_file_obj():
 )
 @pytest.mark.parametrize("index", ["single", "multi", None])
 def test_io_yaml(index):
-    """Test read and write operation on file names."""
+    """Test read and write operation on yaml strings, files and streams."""
     schema = _create_schema(index)
 
     # pass in a file name
@@ -719,6 +720,67 @@ def test_io_yaml(index):
         assert schema_from_yaml == schema
 
 
+@pytest.mark.parametrize("index", ["single", "multi", None])
+def test_io_json(index):
+    """Test read and write operation on json strings, files and streams."""
+    schema = _create_schema(index)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Check Path
+        pth = Path(tmpdir) / "something.json"
+        output = io.to_json(schema, pth)
+        assert output is None
+        schema_from_json = io.from_json(pth)
+        assert schema_from_json == schema
+
+        # DataFrameSchema method
+        output = schema.to_json(pth)
+        assert output is None
+        schema_from_json = schema.from_json(pth)
+        assert schema_from_json == schema
+
+        # Check path as string
+        output = io.to_json(schema, str(pth))
+        assert output is None
+        schema_from_json = io.from_json(str(pth))
+        assert schema_from_json == schema
+
+        # DataFrameSchema method
+        output = schema.to_json(str(pth))
+        assert output is None
+        schema_from_json = schema.from_json(pth)
+        assert schema_from_json == schema
+
+        # Check schema encoded as a string
+        text = io.to_json(schema)
+        assert text is not None
+        assert isinstance(text, str)
+        schema_from_json = io.from_json(text)
+        assert schema_from_json == schema
+
+        # DataFrameSchema method
+        text = schema.to_json()
+        assert text is not None
+        schema_from_json = schema.from_json(text)
+        assert schema_from_json == schema
+
+        # Check schema encoded in a stream
+        stream = StringIO()
+        output = io.to_json(schema, stream)
+        assert output is None
+        stream.seek(0)
+        schema_from_json = io.from_json(stream)
+        assert schema_from_json == schema
+
+        # DataFrameSchema method
+        stream = StringIO()
+        output = schema.to_json(stream)
+        assert output is None
+        stream.seek(0)
+        schema_from_json = schema.from_json(stream)
+        assert schema_from_json == schema
+
+
 @pytest.mark.skipif(
     platform.system() == "Windows",
     reason="skipping due to issues with opening file names for temp files.",