TYP: core.internals (#37058)

Author: jbrockmendel

pandas/core/arrays/base.py

@@ -1081,6 +1081,19 @@ def _formatter(self, boxed: bool = False) -> Callable[[Any], Optional[str]]:
     # Reshaping
     # ------------------------------------------------------------------------
 
+    def transpose(self, *axes) -> "ExtensionArray":
+        """
+        Return a transposed view on this array.
+
+        Because ExtensionArrays are always 1D, this is a no-op.  It is included
+        for compatibility with np.ndarray.
+        """
+        return self[:]
+
+    @property
+    def T(self) -> "ExtensionArray":
+        return self.transpose()
+
     def ravel(self, order="C") -> "ExtensionArray":
         """
         Return a flattened view on this array.

pandas/core/arrays/sparse/array.py

@@ -1310,19 +1310,6 @@ def mean(self, axis=0, *args, **kwargs):
             nsparse = self.sp_index.ngaps
             return (sp_sum + self.fill_value * nsparse) / (ct + nsparse)
 
-    def transpose(self, *axes) -> "SparseArray":
-        """
-        Returns the SparseArray.
-        """
-        return self
-
-    @property
-    def T(self) -> "SparseArray":
-        """
-        Returns the SparseArray.
-        """
-        return self
-
     # ------------------------------------------------------------------------
     # Ufuncs
     # ------------------------------------------------------------------------

pandas/core/internals/blocks.py

@@ -1,7 +1,7 @@
 from datetime import datetime, timedelta
 import inspect
 import re
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional, Type, Union, cast
 import warnings
 
 import numpy as np
@@ -93,6 +93,8 @@ class Block(PandasObject):
     Index-ignorant; let the container take care of that
     """
 
+    values: Union[np.ndarray, ExtensionArray]
+
     __slots__ = ["_mgr_locs", "values", "ndim"]
     is_numeric = False
     is_float = False
@@ -194,7 +196,9 @@ def _consolidate_key(self):
     @property
     def is_view(self) -> bool:
         """ return a boolean if I am possibly a view """
-        return self.values.base is not None
+        values = self.values
+        values = cast(np.ndarray, values)
+        return values.base is not None
 
     @property
     def is_categorical(self) -> bool:
@@ -1465,6 +1469,7 @@ def where_func(cond, values, other):
         result_blocks: List["Block"] = []
         for m in [mask, ~mask]:
             if m.any():
+                result = cast(np.ndarray, result)  # EABlock overrides where
                 taken = result.take(m.nonzero()[0], axis=axis)
                 r = maybe_downcast_numeric(taken, self.dtype)
                 nb = self.make_block(r.T, placement=self.mgr_locs[m])
@@ -1619,6 +1624,8 @@ class ExtensionBlock(Block):
     _validate_ndim = False
     is_extension = True
 
+    values: ExtensionArray
+
     def __init__(self, values, placement, ndim=None):
         """
         Initialize a non-consolidatable block.
@@ -2197,9 +2204,7 @@ def astype(self, dtype, copy: bool = False, errors: str = "raise"):
             if copy:
                 # this should be the only copy
                 values = values.copy()
-            if getattr(values, "tz", None) is None:
-                values = DatetimeArray(values).tz_localize("UTC")
-            values = values.tz_convert(dtype.tz)
+            values = DatetimeArray._simple_new(values.view("i8"), dtype=dtype)
             return self.make_block(values)
 
         # delegate
@@ -2243,6 +2248,8 @@ def set(self, locs, values):
 class DatetimeTZBlock(ExtensionBlock, DatetimeBlock):
     """ implement a datetime64 block with a tz attribute """
 
+    values: DatetimeArray
+
     __slots__ = ()
     is_datetimetz = True
     is_extension = True
@@ -2670,6 +2677,8 @@ def get_block_type(values, dtype=None):
     dtype = dtype or values.dtype
     vtype = dtype.type
 
+    cls: Type[Block]
+
     if is_sparse(dtype):
         # Need this first(ish) so that Sparse[datetime] is sparse
         cls = ExtensionBlock

pandas/core/internals/construction.py

@@ -9,7 +9,7 @@
 import numpy.ma as ma
 
 from pandas._libs import lib
-from pandas._typing import Axis, DtypeObj, Scalar
+from pandas._typing import Axis, DtypeObj, Label, Scalar
 
 from pandas.core.dtypes.cast import (
     construct_1d_arraylike_from_scalar,
@@ -436,7 +436,7 @@ def get_names_from_index(data):
     if not has_some_name:
         return ibase.default_index(len(data))
 
-    index = list(range(len(data)))
+    index: List[Label] = list(range(len(data)))
     count = 0
     for i, s in enumerate(data):
         n = getattr(s, "name", None)

pandas/tests/extension/base/reshaping.py

@@ -333,6 +333,20 @@ def test_ravel(self, data):
         assert data[0] == data[1]
 
     def test_transpose(self, data):
+        result = data.transpose()
+        assert type(result) == type(data)
+
+        # check we get a new object
+        assert result is not data
+
+        # If we ever _did_ support 2D, shape should be reversed
+        assert result.shape == data.shape[::-1]
+
+        # Check that we have a view, not a copy
+        result[0] = result[1]
+        assert data[0] == data[1]
+
+    def test_transpose_frame(self, data):
         df = pd.DataFrame({"A": data[:4], "B": data[:4]}, index=["a", "b", "c", "d"])
         result = df.T
         expected = pd.DataFrame(

pandas/tests/extension/test_numpy.py

@@ -377,8 +377,8 @@ def test_merge_on_extension_array_duplicates(self, data):
         super().test_merge_on_extension_array_duplicates(data)
 
     @skip_nested
-    def test_transpose(self, data):
-        super().test_transpose(data)
+    def test_transpose_frame(self, data):
+        super().test_transpose_frame(data)
 
 
 class TestSetitem(BaseNumPyTests, base.BaseSetitemTests):

pandas/tests/extension/test_sparse.py

@@ -155,6 +155,10 @@ def test_merge(self, data, na_value):
         self._check_unsupported(data)
         super().test_merge(data, na_value)
 
+    @pytest.mark.xfail(reason="SparseArray does not support setitem")
+    def test_transpose(self, data):
+        super().test_transpose(data)
+
 
 class TestGetitem(BaseSparseTests, base.BaseGetitemTests):
     def test_get(self, data):

setup.cfg

@@ -187,12 +187,6 @@ check_untyped_defs=False
 [mypy-pandas.core.indexes.multi]
 check_untyped_defs=False
 
-[mypy-pandas.core.internals.blocks]
-check_untyped_defs=False
-
-[mypy-pandas.core.internals.construction]
-check_untyped_defs=False
-
 [mypy-pandas.core.resample]
 check_untyped_defs=False