diff --git a/hamilton/plugins/polars_post_1_0_0_extensions.py b/hamilton/plugins/polars_post_1_0_0_extensions.py index dab1a3973..3df25674e 100644 --- a/hamilton/plugins/polars_post_1_0_0_extensions.py +++ b/hamilton/plugins/polars_post_1_0_0_extensions.py @@ -523,8 +523,9 @@ def _get_loading_kwargs(self): def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]: df = pl.read_json(self.source, **self._get_loading_kwargs()) - metadata = utils.get_file_metadata(self.source) - return df, metadata + if isinstance(self.source, (str, Path)): + return df, utils.get_file_metadata(self.source) + return df, utils.get_dataframe_metadata(df) @classmethod def name(cls) -> str: @@ -549,7 +550,9 @@ def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: data = data.collect() data.write_json(self.file) - return utils.get_file_and_dataframe_metadata(self.file, data) + if isinstance(self.file, (str, Path)): + return utils.get_file_and_dataframe_metadata(self.file, data) + return utils.get_dataframe_metadata(data) @classmethod def name(cls) -> str: @@ -581,8 +584,9 @@ def _get_loading_kwargs(self): def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]: df = pl.read_ndjson(self.source, **self._get_loading_kwargs()) - metadata = utils.get_file_metadata(self.source) - return df, metadata + if isinstance(self.source, (str, Path)): + return df, utils.get_file_metadata(self.source) + return df, utils.get_dataframe_metadata(df) @classmethod def name(cls) -> str: @@ -607,7 +611,9 @@ def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: data = data.collect() data.write_ndjson(self.file) - return utils.get_file_and_dataframe_metadata(self.file, data) + if isinstance(self.file, (str, Path)): + return utils.get_file_and_dataframe_metadata(self.file, data) + return utils.get_dataframe_metadata(data) @classmethod def name(cls) -> str: diff --git a/tests/plugins/test_polars_extensions.py b/tests/plugins/test_polars_extensions.py index 2709dd85c..967571141 100644 --- a/tests/plugins/test_polars_extensions.py +++ b/tests/plugins/test_polars_extensions.py @@ -147,6 +147,24 @@ def test_polars_ndjson(df: pl.DataFrame, tmp_path: pathlib.Path) -> None: polars.testing.assert_frame_equal(df, df2) +@pytest.mark.parametrize( + ("writer_cls", "reader_cls"), + [(PolarsJSONWriter, PolarsJSONReader), (PolarsNDJSONWriter, PolarsNDJSONReader)], +) +def test_polars_json_file_objects(df: pl.DataFrame, writer_cls, reader_cls) -> None: + buffer = io.BytesIO() + write_metadata = writer_cls(file=buffer).save_data(df) + + df2, read_metadata = reader_cls(source=io.BytesIO(buffer.getvalue())).load_data(pl.DataFrame) + _, bytes_metadata = reader_cls(source=buffer.getvalue()).load_data(pl.DataFrame) + + polars.testing.assert_frame_equal(df, df2) + for metadata in (write_metadata, read_metadata, bytes_metadata): + assert "file_metadata" not in metadata + assert metadata["dataframe_metadata"]["rows"] == 2 + assert metadata["dataframe_metadata"]["column_names"] == ["a", "b"] + + def test_polars_avro(df: pl.DataFrame, tmp_path: pathlib.Path) -> None: file = tmp_path / "test.avro"