diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py index ba8c211fec9e..f810028676e1 100644 --- 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 index ec531dbd8ca2..77d005823c6f 100644 --- 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 """ - 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 index 46c4cd561e00..58e743b173c9 100644 --- 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,