Enhancment/1886 Add json_normalize to pandas read formats (unionai-os…

…s#1892) * add json_normalize to pandas read formats * add unit test * format code with black/fix unit test
max-raphael · Jan 24, 2025 · aca3850 · aca3850
1 parent 3b638df
commit aca3850
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
diff --git a/pandera/typing/formats.py b/pandera/typing/formats.py
@@ -42,6 +42,9 @@ class Formats(Enum):
     #: python pickle file format
     pickle = "pickle"
 
+    #: python json_normalize
+    json_normalize = "json_normalize"
+
 
 Format = Union[
     Literal[Formats.csv],
@@ -50,4 +53,5 @@ class Formats(Enum):
     Literal[Formats.feather],
     Literal[Formats.parquet],
     Literal[Formats.pickle],
+    Literal[Formats.json_normalize],
 ]
diff --git a/pandera/typing/pandas.py b/pandera/typing/pandas.py
@@ -122,6 +122,7 @@ def from_format(cls, obj: Any, config) -> pd.DataFrame:
                 Formats.feather: pd.read_feather,
                 Formats.parquet: pd.read_parquet,
                 Formats.pickle: pd.read_pickle,
+                Formats.json_normalize: pd.json_normalize,
             }[Formats(config.from_format)]
 
         return reader(obj, **(config.from_format_kwargs or {}))  # type: ignore

diff --git a/tests/core/test_from_to_format_conversions.py b/tests/core/test_from_to_format_conversions.py
@@ -64,6 +64,11 @@ class Config:
         from_format = pd.read_pickle
 
 
+class InSchemaJsonNormalize(InSchema):
+    class Config:
+        from_format = "json_normalize"
+
+
 class OutSchema(InSchema):
     float_col: pa.typing.Series[float]
 
@@ -194,6 +199,7 @@ def _needs_pyarrow(schema) -> bool:
         [InSchemaParquet, lambda df, x: df.to_parquet(x), io.BytesIO],
         [InSchemaPickle, lambda df, x: df.to_pickle(x), io.BytesIO],
         [InSchemaPickleCallable, lambda df, x: df.to_pickle(x), io.BytesIO],
+        [InSchemaJsonNormalize, lambda df: df.to_dict(orient="records"), None],
     ],
 )
 def test_from_format(schema, to_fn, buf_cls):