TYP: _json.py (#56539)

* TYP: _json.py

* fix index special-casing

Author: jbrockmendel

pandas/io/json/_json.py

@@ -14,6 +14,7 @@
     Generic,
     Literal,
     TypeVar,
+    final,
     overload,
 )
 import warnings
@@ -37,11 +38,11 @@
     is_string_dtype,
 )
 from pandas.core.dtypes.dtypes import PeriodDtype
-from pandas.core.dtypes.generic import ABCIndex
 
 from pandas import (
     ArrowDtype,
     DataFrame,
+    Index,
     MultiIndex,
     Series,
     isna,
@@ -1119,10 +1120,11 @@ class Parser:
         "us": 31536000000000,
         "ns": 31536000000000000,
     }
+    json: str
 
     def __init__(
         self,
-        json,
+        json: str,
         orient,
         dtype: DtypeArg | None = None,
         convert_axes: bool = True,
@@ -1157,7 +1159,8 @@ def __init__(
         self.obj: DataFrame | Series | None = None
         self.dtype_backend = dtype_backend
 
-    def check_keys_split(self, decoded) -> None:
+    @final
+    def check_keys_split(self, decoded: dict) -> None:
         """
         Checks that dict has only the appropriate keys for orient='split'.
         """
@@ -1166,6 +1169,7 @@ def check_keys_split(self, decoded) -> None:
             bad_keys_joined = ", ".join(bad_keys)
             raise ValueError(f"JSON data had unexpected key(s): {bad_keys_joined}")
 
+    @final
     def parse(self):
         self._parse()
 
@@ -1179,34 +1183,41 @@ def parse(self):
     def _parse(self):
         raise AbstractMethodError(self)
 
+    @final
     def _convert_axes(self) -> None:
         """
         Try to convert axes.
         """
         obj = self.obj
         assert obj is not None  # for mypy
         for axis_name in obj._AXIS_ORDERS:
-            new_axis, result = self._try_convert_data(
+            ax = obj._get_axis(axis_name)
+            ser = Series(ax, dtype=ax.dtype, copy=False)
+            new_ser, result = self._try_convert_data(
                 name=axis_name,
-                data=obj._get_axis(axis_name),
+                data=ser,
                 use_dtypes=False,
                 convert_dates=True,
+                is_axis=True,
             )
             if result:
+                new_axis = Index(new_ser, dtype=new_ser.dtype, copy=False)
                 setattr(self.obj, axis_name, new_axis)
 
     def _try_convert_types(self):
         raise AbstractMethodError(self)
 
+    @final
     def _try_convert_data(
         self,
         name: Hashable,
-        data,
+        data: Series,
         use_dtypes: bool = True,
         convert_dates: bool | list[str] = True,
-    ):
+        is_axis: bool = False,
+    ) -> tuple[Series, bool]:
         """
-        Try to parse a ndarray like into a column by inferring dtype.
+        Try to parse a Series into a column by inferring dtype.
         """
         # don't try to coerce, unless a force conversion
         if use_dtypes:
@@ -1242,7 +1253,7 @@ def _try_convert_data(
             if result:
                 return new_data, True
 
-        if self.dtype_backend is not lib.no_default and not isinstance(data, ABCIndex):
+        if self.dtype_backend is not lib.no_default and not is_axis:
             # Fall through for conversion later on
             return data, True
         elif is_string_dtype(data.dtype):
@@ -1285,7 +1296,8 @@ def _try_convert_data(
 
         return data, True
 
-    def _try_convert_to_date(self, data):
+    @final
+    def _try_convert_to_date(self, data: Series) -> tuple[Series, bool]:
         """
         Try to parse a ndarray like into a date column.
 
@@ -1335,13 +1347,11 @@ def _try_convert_to_date(self, data):
             return new_data, True
         return data, False
 
-    def _try_convert_dates(self):
-        raise AbstractMethodError(self)
-
 
 class SeriesParser(Parser):
     _default_orient = "index"
     _split_keys = ("name", "index", "data")
+    obj: Series | None
 
     def _parse(self) -> None:
         data = ujson_loads(self.json, precise_float=self.precise_float)
@@ -1366,6 +1376,7 @@ def _try_convert_types(self) -> None:
 class FrameParser(Parser):
     _default_orient = "columns"
     _split_keys = ("columns", "index", "data")
+    obj: DataFrame | None
 
     def _parse(self) -> None:
         json = self.json
@@ -1403,20 +1414,24 @@ def _parse(self) -> None:
                 ujson_loads(json, precise_float=self.precise_float), dtype=None
             )
 
-    def _process_converter(self, f, filt=None) -> None:
+    def _process_converter(
+        self,
+        f: Callable[[Hashable, Series], tuple[Series, bool]],
+        filt: Callable[[Hashable], bool] | None = None,
+    ) -> None:
         """
         Take a conversion function and possibly recreate the frame.
         """
         if filt is None:
-            filt = lambda col, c: True
+            filt = lambda col: True
 
         obj = self.obj
         assert obj is not None  # for mypy
 
         needs_new_obj = False
         new_obj = {}
         for i, (col, c) in enumerate(obj.items()):
-            if filt(col, c):
+            if filt(col):
                 new_data, result = f(col, c)
                 if result:
                     c = new_data
@@ -1453,6 +1468,10 @@ def is_ok(col) -> bool:
             """
             Return if this col is ok to try for a date parse.
             """
+            if col in convert_dates:
+                return True
+            if not self.keep_default_dates:
+                return False
             if not isinstance(col, str):
                 return False
 
@@ -1467,9 +1486,4 @@ def is_ok(col) -> bool:
                 return True
             return False
 
-        self._process_converter(
-            lambda col, c: self._try_convert_to_date(c),
-            lambda col, c: (
-                (self.keep_default_dates and is_ok(col)) or col in convert_dates
-            ),
-        )
+        self._process_converter(lambda col, c: self._try_convert_to_date(c), filt=is_ok)