TYP: check_untyped_defs pandas.core.computation.align (#30550)

Author: simonjayhawkins

pandas/core/computation/align.py

@@ -2,10 +2,12 @@
 """
 
 from functools import partial, wraps
+from typing import Dict, Optional, Sequence, Tuple, Type, Union
 import warnings
 
 import numpy as np
 
+from pandas._typing import FrameOrSeries
 from pandas.errors import PerformanceWarning
 
 from pandas.core.dtypes.generic import ABCDataFrame, ABCSeries
@@ -15,22 +17,27 @@
 from pandas.core.computation.common import result_type_many
 
 
-def _align_core_single_unary_op(term):
+def _align_core_single_unary_op(
+    term,
+) -> Tuple[Union[partial, Type[FrameOrSeries]], Optional[Dict[str, int]]]:
+
+    typ: Union[partial, Type[FrameOrSeries]]
+    axes: Optional[Dict[str, int]] = None
+
     if isinstance(term.value, np.ndarray):
         typ = partial(np.asanyarray, dtype=term.value.dtype)
     else:
         typ = type(term.value)
-    ret = (typ,)
+        if hasattr(term.value, "axes"):
+            axes = _zip_axes_from_type(typ, term.value.axes)
 
-    if not hasattr(term.value, "axes"):
-        ret += (None,)
-    else:
-        ret += (_zip_axes_from_type(typ, term.value.axes),)
-    return ret
+    return typ, axes
 
 
-def _zip_axes_from_type(typ, new_axes):
-    axes = {ax_name: new_axes[ax_ind] for ax_ind, ax_name in typ._AXIS_NAMES.items()}
+def _zip_axes_from_type(
+    typ: Type[FrameOrSeries], new_axes: Sequence[int]
+) -> Dict[str, int]:
+    axes = {name: new_axes[i] for i, name in typ._AXIS_NAMES.items()}
     return axes
 
 

setup.cfg

@@ -163,9 +163,6 @@ check_untyped_defs=False
 [mypy-pandas.core.base]
 check_untyped_defs=False
 
-[mypy-pandas.core.computation.align]
-check_untyped_defs=False
-
 [mypy-pandas.core.computation.expr]
 check_untyped_defs=False