Separate write_json from to_json (CQS)

Author: nathanjmcdougall

pins/_adaptors.py

@@ -24,21 +24,14 @@ class Adaptor:
     def __init__(self, data: Any) -> None:
         self._d = data
 
-    @overload
-    def write_json(self, file: str) -> None: ...
-    @overload
-    def write_json(self, file: None = ...) -> str: ...
-    def write_json(self, file: str | None = None) -> str | None:
-        if file is None:
-            msg = (
-                f"Writing to JSON string rather than file is not supported for "
-                f"{type(self._d)}"
-            )
-            raise NotImplementedError(msg)
+    def write_json(self, file: str) -> None:
+        with open(file, "w") as f:
+            f.write(self.to_json())
 
+    def to_json(self) -> str:
         import json
 
-        json.dump(self._d, open(file, mode="w"))
+        return json.dumps(self._d)
 
     def write_joblib(self, file: str) -> None:
         import joblib
@@ -101,7 +94,7 @@ def head(self, n: int) -> DFAdaptor: ...
     def data_preview(self) -> str:
         # TODO(compat) is 100 hard-coded?
         # Note that we go df -> json -> dict, to take advantage of type conversions in the dataframe library
-        data: list[dict[Any, Any]] = json.loads(self.head(100).write_json())
+        data: list[dict[Any, Any]] = json.loads(self.head(100).to_json())
         columns = [
             {"name": [col], "label": [col], "align": ["left"], "type": [""]}
             for col in self.columns
@@ -135,18 +128,7 @@ def shape(self) -> tuple[int, int]:
     def head(self, n: int) -> PandasAdaptor:
         return PandasAdaptor(self._d.head(n))
 
-    @overload
-    def write_json(self, file: str) -> None: ...
-    @overload
-    def write_json(self, file: None) -> str: ...
-    def write_json(self, file: str | None = None) -> str | None:
-        if file is not None:
-            msg = (
-                f"Writing to file rather than JSON string is not supported for "
-                f"{type(self._d)}"
-            )
-            raise NotImplementedError(msg)
-
+    def to_json(self) -> str:
         return self._d.to_json(orient="records")
 
     def write_csv(self, file: str) -> None:

pins/tests/test_adaptors.py

@@ -107,11 +107,11 @@ def test_head(self):
         assert isinstance(adaptor.head(1), PandasAdaptor)
         assert_frame_equal(adaptor.head(1)._d, expected._d)
 
-    def test_write_json(self):
+    def test_to_json(self):
         df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
         adaptor = PandasAdaptor(df)
         assert isinstance(adaptor, DFAdaptor)
-        assert adaptor.write_json() == """[{"a":1,"b":4},{"a":2,"b":5},{"a":3,"b":6}]"""
+        assert adaptor.to_json() == """[{"a":1,"b":4},{"a":2,"b":5},{"a":3,"b":6}]"""
 
     def test_write_csv(self, tmp_path: Path):
         df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})