Add NarwhalsFilePublisher (#93)

* Add NarwhalsFilePublisher.
* Improve consistency with rest of the code. Use narwhals.stable.v1 and use exttypes.narwhals.
* Require eager frame in NarwhalsFilePublisher.

Signed-off-by: Hin Tse <5867507+hintse@users.noreply.github.com>

Author: hintse

ccflow/publishers/file.py

@@ -1,13 +1,15 @@
 import pickle
 from typing import IO, Any, Callable, Dict, Generic
 
+import narwhals.stable.v1 as nw
 import pandas as pd
 import yaml
 from cloudpathlib import AnyPath
 from pydantic import Field, field_validator
 from typing_extensions import Literal, override
 
 from ..exttypes import JinjaTemplate
+from ..exttypes.narwhals import DataFrameT
 from ..publisher import BasePublisher
 from ..serialization import orjson_dumps
 from ..utils import PydanticDictOptions, PydanticModelType, dict_to_model
@@ -16,6 +18,7 @@
     "DictTemplateFilePublisher",
     "GenericFilePublisher",
     "JSONPublisher",
+    "NarwhalsFilePublisher",
     "PandasFilePublisher",
     "PicklePublisher",
     "PydanticJSONPublisher",
@@ -190,3 +193,25 @@ def __call__(self) -> AnyPath:
             mode=self.mode,
             kwargs=self.kwargs,
         )()
+
+
+class NarwhalsFilePublisher(BasePublisher):
+    """Publish a narwhals data frame to a file using an appropriate method on nw.DataFrame."""
+
+    data: DataFrameT = None
+    kwargs: Dict[str, Any] = Field(default_factory=dict)
+    func: str = "write_csv"  # The access function must be able to write to a buffer or file-like object.
+    suffix: str = ".csv"
+    mode: Literal["w", "wb"] = "w"
+
+    @override
+    def __call__(self) -> AnyPath:
+        return GenericFilePublisher(
+            name=self.name,
+            name_params=self.name_params,
+            data=self.data,
+            dump=getattr(nw.DataFrame, self.func),
+            suffix=self.suffix,
+            mode=self.mode,
+            kwargs=self.kwargs,
+        )()

ccflow/tests/publishers/test_file.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 from unittest import TestCase
 
+import narwhals.stable.v1 as nw
 import pandas as pd
 from pydantic import BaseModel as PydanticBaseModel
 
@@ -13,6 +14,7 @@
     DictTemplateFilePublisher,
     GenericFilePublisher,
     JSONPublisher,
+    NarwhalsFilePublisher,
     PandasFilePublisher,
     PicklePublisher,
     YAMLPublisher,
@@ -162,3 +164,36 @@ def test_pandas_feather(self):
             self.assertEqual(path, Path("test_pandas.f"))
             df = pd.read_feather(path)
             pd.testing.assert_frame_equal(df, p.data)
+
+    def test_narwhals_csv(self):
+        with tempfile.TemporaryDirectory() as tempdir:
+            os.chdir(tempdir)
+            p = NarwhalsFilePublisher(
+                name="test_{{param}}",
+                name_params={"param": "narwhals"},
+                func="write_csv",
+                suffix=".csv",
+            )
+            df = pd.DataFrame({"a": [1, 2, 3], "b": ["foo", "bar", "baz"]})
+            p.data = nw.from_native(df)
+            path = p()
+            self.assertEqual(path, Path("test_narwhals.csv"))
+            df2 = pd.read_csv(path)
+            pd.testing.assert_frame_equal(df, df2)
+
+    def test_narwhals_parquet(self):
+        with tempfile.TemporaryDirectory() as tempdir:
+            os.chdir(tempdir)
+            p = NarwhalsFilePublisher(
+                name="test_{{param}}",
+                name_params={"param": "narwhals"},
+                func="write_parquet",
+                suffix=".parquet",
+                mode="wb",
+            )
+            df = pd.DataFrame({"a": [1, 2, 3], "b": ["foo", "bar", "baz"]})
+            p.data = nw.from_native(df)
+            path = p()
+            self.assertEqual(path, Path("test_narwhals.parquet"))
+            df2 = pd.read_parquet(path)
+            pd.testing.assert_frame_equal(df, df2)

docs/wiki/Key-Features.md

@@ -58,6 +58,7 @@ The following table summarizes the "publisher" models.
 | `GenericFilePublisher`       | `ccflow.publishers` | Publish data using a generic "dump" Callable. Uses `smart_open` under the hood so that local and cloud paths are supported.                                 |
 | `JSONPublisher`              | `ccflow.publishers` | Publish data to file in JSON format.                                                                                                                        |
 | `PandasFilePublisher`        | `ccflow.publishers` | Publish a pandas data frame to a file using an appropriate method on pd.DataFrame. For large-scale exporting (using parquet), see `PandasParquetPublisher`. |
+| `NarwhalsFilePublisher`      | `ccflow.publishers` | Publish a narwhals data frame to a file using an appropriate method on nw.DataFrame.                                                                        |
 | `PicklePublisher`            | `ccflow.publishers` | Publish data to a pickle file.                                                                                                                              |
 | `PydanticJSONPublisher`      | `ccflow.publishers` | Publish a pydantic model to a json file. See [Pydantic modeljson](https://docs.pydantic.dev/latest/concepts/serialization/#modelmodel_dump)                 |
 | `YAMLPublisher`              | `ccflow.publishers` | Publish data to file in YAML format.                                                                                                                        |