From 6f872be229b107c27230914c5f004622a91b10d2 Mon Sep 17 00:00:00 2001
From: digit50 <93092234+digit50@users.noreply.github.com>
Date: Thu, 15 Oct 2026 16:21:32 +0300
Subject: [PATCH] Support file objects in Polars JSON and NDJSON readers and
 writers

The Polars JSON and NDJSON loaders and savers accept file objects and bytes, but passed them to get_file_metadata, which only handles paths, so a BytesIO raised AttributeError (on the writer after the data was written). File metadata is now computed only for str and Path, and dataframe metadata is returned for file objects and bytes.

Closes https://github.com/apache/hamilton/issues/1758
---
 .../plugins/polars_post_1_0_0_extensions.py    | 18 ++++++++++++------
 tests/plugins/test_polars_extensions.py        | 18 ++++++++++++++++++
 2 files changed, 30 insertions(+), 6 deletions(-)

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"