Reuse improved parsing logic

pola-rs · May 28, 2024 · 816c70e · 816c70e
1 parent db475c5
commit 816c70e
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 31 deletions.
diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py
@@ -2347,22 +2347,26 @@ def write_ndjson(self, file: IOBase | str | Path | None = None) -> str | None:
         >>> df.write_ndjson()
         '{"foo":1,"bar":6}\n{"foo":2,"bar":7}\n{"foo":3,"bar":8}\n'
         """
-        if isinstance(file, (str, Path)):
-            file = normalize_filepath(file)
-        to_string_io = (file is not None) and isinstance(file, StringIO)
-        if file is None or to_string_io:
+
+        def write_ndjson_to_string() -> str:
             with BytesIO() as buf:
                 self._df.write_ndjson(buf)
-                json_bytes = buf.getvalue()
+                ndjson_bytes = buf.getvalue()
+            return ndjson_bytes.decode("utf8")
 
-            json_str = json_bytes.decode("utf8")
-            if to_string_io:
-                file.write(json_str)  # type: ignore[union-attr]
-            else:
-                return json_str
+        if file is None:
+            return write_ndjson_to_string()
+        elif isinstance(file, StringIO):
+            ndjson_str = write_ndjson_to_string()
+            file.write(ndjson_str)
+            return None
+        elif isinstance(file, (str, Path)):
+            file = normalize_filepath(file)
+            self._df.write_ndjson(file)
+            return None
         else:
             self._df.write_ndjson(file)
-        return None
+            return None
 
     @overload
     def write_csv(

diff --git a/py-polars/polars/expr/meta.py b/py-polars/polars/expr/meta.py
@@ -291,22 +291,26 @@ def serialize(self, file: IOBase | str | Path | None = None) -> str | None:
         >>> pl.Expr.deserialize(StringIO(json))  # doctest: +ELLIPSIS
         <Expr ['col("foo").sum().over([col("ba…'] at ...>
         """
-        if isinstance(file, (str, Path)):
-            file = normalize_filepath(file)
-        to_string_io = (file is not None) and isinstance(file, StringIO)
-        if file is None or to_string_io:
+
+        def serialize_to_string() -> str:
             with BytesIO() as buf:
                 self._pyexpr.serialize(buf)
                 json_bytes = buf.getvalue()
+            return json_bytes.decode("utf8")
 
-            json_str = json_bytes.decode("utf8")
-            if to_string_io:
-                file.write(json_str)  # type: ignore[union-attr]
-            else:
-                return json_str
+        if file is None:
+            return serialize_to_string()
+        elif isinstance(file, StringIO):
+            json_str = serialize_to_string()
+            file.write(json_str)
+            return None
+        elif isinstance(file, (str, Path)):
+            file = normalize_filepath(file)
+            self._pyexpr.serialize(file)
+            return None
         else:
             self._pyexpr.serialize(file)
-        return None
+            return None
 
     @overload
     def write_json(self, file: None = ...) -> str: ...

diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py
@@ -589,22 +589,26 @@ def serialize(self, file: IOBase | str | Path | None = None) -> str | None:
         │ 6   │
         └─────┘
         """
-        if isinstance(file, (str, Path)):
-            file = normalize_filepath(file)
-        to_string_io = (file is not None) and isinstance(file, StringIO)
-        if file is None or to_string_io:
+
+        def serialize_to_string() -> str:
             with BytesIO() as buf:
                 self._ldf.serialize(buf)
                 json_bytes = buf.getvalue()
+            return json_bytes.decode("utf8")
 
-            json_str = json_bytes.decode("utf8")
-            if to_string_io:
-                file.write(json_str)  # type: ignore[union-attr]
-            else:
-                return json_str
+        if file is None:
+            return serialize_to_string()
+        elif isinstance(file, StringIO):
+            json_str = serialize_to_string()
+            file.write(json_str)
+            return None
+        elif isinstance(file, (str, Path)):
+            file = normalize_filepath(file)
+            self._ldf.serialize(file)
+            return None
         else:
             self._ldf.serialize(file)
-        return None
+            return None
 
     def pipe(
         self,