check_untyped_defs pandas.core.computation.align

Author: simonjayhawkins

pandas/core/computation/align.py

@@ -2,6 +2,7 @@
 """
 
 from functools import partial, wraps
+from typing import Dict, Optional, Sequence, Tuple, Type, Union
 import warnings
 
 import numpy as np
@@ -14,22 +15,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[pd.core.generic.NDFrame]], Optional[Dict[str, int]]]:
+
+    typ: Union[partial, Type[pd.core.generic.NDFrame]]
+    axes: Optional[Dict] = 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[pd.core.generic.NDFrame], new_axes: Sequence[int]
+) -> Dict[str, int]:
+    axes = {name: new_axes[i] for i, name in typ._AXIS_NAMES.items()}
     return axes
 
 

pandas/core/generic.py

@@ -171,6 +171,7 @@ class NDFrame(PandasObject, SelectionMixin):
     _metadata = []  # type: List[str]
     _is_copy = None
     _data = None  # type: BlockManager
+    _AXIS_NAMES: Dict[int, str]
 
     # ----------------------------------------------------------------------
     # Constructors

setup.cfg

@@ -185,9 +185,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