REF: eliminate method _write() in json writers (#36218)

Author: ivanovmg

pandas/io/json/_json.py

@@ -1,22 +1,29 @@
+from abc import ABC, abstractmethod
 from collections import abc
 import functools
 from io import BytesIO, StringIO
 from itertools import islice
 import os
-from typing import IO, Any, Callable, List, Optional, Tuple, Type
+from typing import IO, Any, Callable, List, Mapping, Optional, Tuple, Type, Union
 
 import numpy as np
 
 import pandas._libs.json as json
 from pandas._libs.tslibs import iNaT
-from pandas._typing import CompressionOptions, JSONSerializable, StorageOptions
+from pandas._typing import (
+    CompressionOptions,
+    IndexLabel,
+    JSONSerializable,
+    StorageOptions,
+)
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import deprecate_kwarg, deprecate_nonkeyword_arguments
 
 from pandas.core.dtypes.common import ensure_str, is_period_dtype
 
 from pandas import DataFrame, MultiIndex, Series, isna, to_datetime
 from pandas.core.construction import create_series_with_explicit_dtype
+from pandas.core.generic import NDFrame
 from pandas.core.reshape.concat import concat
 
 from pandas.io.common import get_compression_method, get_filepath_or_buffer, get_handle
@@ -33,7 +40,7 @@
 # interface to/from
 def to_json(
     path_or_buf,
-    obj,
+    obj: NDFrame,
     orient: Optional[str] = None,
     date_format: str = "epoch",
     double_precision: int = 10,
@@ -110,7 +117,7 @@ def to_json(
             path_or_buf.close()
 
 
-class Writer:
+class Writer(ABC):
     _default_orient: str
 
     def __init__(
@@ -146,75 +153,52 @@ def _format_axes(self):
         raise AbstractMethodError(self)
 
     def write(self):
-        return self._write(
-            self.obj,
-            self.orient,
-            self.double_precision,
-            self.ensure_ascii,
-            self.date_unit,
-            self.date_format == "iso",
-            self.default_handler,
-            self.indent,
-        )
-
-    def _write(
-        self,
-        obj,
-        orient: Optional[str],
-        double_precision: int,
-        ensure_ascii: bool,
-        date_unit: str,
-        iso_dates: bool,
-        default_handler: Optional[Callable[[Any], JSONSerializable]],
-        indent: int,
-    ):
+        iso_dates = self.date_format == "iso"
         return dumps(
-            obj,
-            orient=orient,
-            double_precision=double_precision,
-            ensure_ascii=ensure_ascii,
-            date_unit=date_unit,
+            self.obj_to_write,
+            orient=self.orient,
+            double_precision=self.double_precision,
+            ensure_ascii=self.ensure_ascii,
+            date_unit=self.date_unit,
             iso_dates=iso_dates,
-            default_handler=default_handler,
-            indent=indent,
+            default_handler=self.default_handler,
+            indent=self.indent,
         )
 
+    @property
+    @abstractmethod
+    def obj_to_write(self) -> Union[NDFrame, Mapping[IndexLabel, Any]]:
+        """Object to write in JSON format."""
+        pass
+
 
 class SeriesWriter(Writer):
     _default_orient = "index"
 
+    @property
+    def obj_to_write(self) -> Union[NDFrame, Mapping[IndexLabel, Any]]:
+        if not self.index and self.orient == "split":
+            return {"name": self.obj.name, "data": self.obj.values}
+        else:
+            return self.obj
+
     def _format_axes(self):
         if not self.obj.index.is_unique and self.orient == "index":
             raise ValueError(f"Series index must be unique for orient='{self.orient}'")
 
-    def _write(
-        self,
-        obj,
-        orient: Optional[str],
-        double_precision: int,
-        ensure_ascii: bool,
-        date_unit: str,
-        iso_dates: bool,
-        default_handler: Optional[Callable[[Any], JSONSerializable]],
-        indent: int,
-    ):
-        if not self.index and orient == "split":
-            obj = {"name": obj.name, "data": obj.values}
-        return super()._write(
-            obj,
-            orient,
-            double_precision,
-            ensure_ascii,
-            date_unit,
-            iso_dates,
-            default_handler,
-            indent,
-        )
-
 
 class FrameWriter(Writer):
     _default_orient = "columns"
 
+    @property
+    def obj_to_write(self) -> Union[NDFrame, Mapping[IndexLabel, Any]]:
+        if not self.index and self.orient == "split":
+            obj_to_write = self.obj.to_dict(orient="split")
+            del obj_to_write["index"]
+        else:
+            obj_to_write = self.obj
+        return obj_to_write
+
     def _format_axes(self):
         """
         Try to format axes if they are datelike.
@@ -232,31 +216,6 @@ def _format_axes(self):
                 f"DataFrame columns must be unique for orient='{self.orient}'."
             )
 
-    def _write(
-        self,
-        obj,
-        orient: Optional[str],
-        double_precision: int,
-        ensure_ascii: bool,
-        date_unit: str,
-        iso_dates: bool,
-        default_handler: Optional[Callable[[Any], JSONSerializable]],
-        indent: int,
-    ):
-        if not self.index and orient == "split":
-            obj = obj.to_dict(orient="split")
-            del obj["index"]
-        return super()._write(
-            obj,
-            orient,
-            double_precision,
-            ensure_ascii,
-            date_unit,
-            iso_dates,
-            default_handler,
-            indent,
-        )
-
 
 class JSONTableWriter(FrameWriter):
     _default_orient = "records"
@@ -331,30 +290,9 @@ def __init__(
         self.orient = "records"
         self.index = index
 
-    def _write(
-        self,
-        obj,
-        orient,
-        double_precision,
-        ensure_ascii,
-        date_unit,
-        iso_dates,
-        default_handler,
-        indent,
-    ):
-        table_obj = {"schema": self.schema, "data": obj}
-        serialized = super()._write(
-            table_obj,
-            orient,
-            double_precision,
-            ensure_ascii,
-            date_unit,
-            iso_dates,
-            default_handler,
-            indent,
-        )
-
-        return serialized
+    @property
+    def obj_to_write(self) -> Union[NDFrame, Mapping[IndexLabel, Any]]:
+        return {"schema": self.schema, "data": self.obj}
 
 
 @deprecate_kwarg(old_arg_name="numpy", new_arg_name=None)