TYP: annotate (pandas-dev#32730)

Author: jbrockmendel

pandas/core/arrays/base.py

@@ -21,7 +21,7 @@
 
 from pandas.core.dtypes.common import is_array_like, is_list_like
 from pandas.core.dtypes.dtypes import ExtensionDtype
-from pandas.core.dtypes.generic import ABCExtensionArray, ABCIndexClass, ABCSeries
+from pandas.core.dtypes.generic import ABCIndexClass, ABCSeries
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import ops
@@ -591,7 +591,7 @@ def dropna(self):
         """
         return self[~self.isna()]
 
-    def shift(self, periods: int = 1, fill_value: object = None) -> ABCExtensionArray:
+    def shift(self, periods: int = 1, fill_value: object = None) -> "ExtensionArray":
         """
         Shift values by desired number.
 
@@ -728,7 +728,7 @@ def _values_for_factorize(self) -> Tuple[np.ndarray, Any]:
         """
         return self.astype(object), np.nan
 
-    def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, ABCExtensionArray]:
+    def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, "ExtensionArray"]:
         """
         Encode the extension array as an enumerated type.
 
@@ -833,7 +833,7 @@ def repeat(self, repeats, axis=None):
 
     def take(
         self, indices: Sequence[int], allow_fill: bool = False, fill_value: Any = None
-    ) -> ABCExtensionArray:
+    ) -> "ExtensionArray":
         """
         Take elements from an array.
 
@@ -922,7 +922,7 @@ def take(self, indices, allow_fill=False, fill_value=None):
         # pandas.api.extensions.take
         raise AbstractMethodError(self)
 
-    def copy(self) -> ABCExtensionArray:
+    def copy(self) -> "ExtensionArray":
         """
         Return a copy of the array.
 
@@ -932,7 +932,7 @@ def copy(self) -> ABCExtensionArray:
         """
         raise AbstractMethodError(self)
 
-    def view(self, dtype=None) -> Union[ABCExtensionArray, np.ndarray]:
+    def view(self, dtype=None) -> ArrayLike:
         """
         Return a view on the array.
 
@@ -943,8 +943,8 @@ def view(self, dtype=None) -> Union[ABCExtensionArray, np.ndarray]:
 
         Returns
         -------
-        ExtensionArray
-            A view of the :class:`ExtensionArray`.
+        ExtensionArray or np.ndarray
+            A view on the :class:`ExtensionArray`'s data.
         """
         # NB:
         # - This must return a *new* object referencing the same data, not self.
@@ -1002,7 +1002,7 @@ def _formatter(self, boxed: bool = False) -> Callable[[Any], Optional[str]]:
     # Reshaping
     # ------------------------------------------------------------------------
 
-    def ravel(self, order="C") -> ABCExtensionArray:
+    def ravel(self, order="C") -> "ExtensionArray":
         """
         Return a flattened view on this array.
 
@@ -1023,8 +1023,8 @@ def ravel(self, order="C") -> ABCExtensionArray:
 
     @classmethod
     def _concat_same_type(
-        cls, to_concat: Sequence[ABCExtensionArray]
-    ) -> ABCExtensionArray:
+        cls, to_concat: Sequence["ExtensionArray"]
+    ) -> "ExtensionArray":
         """
         Concatenate multiple array.
 

pandas/core/arrays/datetimelike.py

@@ -131,7 +131,7 @@ class AttributesMixin:
     _data: np.ndarray
 
     @classmethod
-    def _simple_new(cls, values, **kwargs):
+    def _simple_new(cls, values: np.ndarray, **kwargs):
         raise AbstractMethodError(cls)
 
     @property

pandas/core/arrays/period.py

@@ -31,13 +31,7 @@
     pandas_dtype,
 )
 from pandas.core.dtypes.dtypes import PeriodDtype
-from pandas.core.dtypes.generic import (
-    ABCIndexClass,
-    ABCPeriod,
-    ABCPeriodArray,
-    ABCPeriodIndex,
-    ABCSeries,
-)
+from pandas.core.dtypes.generic import ABCIndexClass, ABCPeriodIndex, ABCSeries
 from pandas.core.dtypes.missing import isna, notna
 
 import pandas.core.algorithms as algos
@@ -48,7 +42,7 @@
 from pandas.tseries.offsets import DateOffset, Tick, _delta_to_tick
 
 
-def _field_accessor(name, alias, docstring=None):
+def _field_accessor(name: str, alias: int, docstring=None):
     def f(self):
         base, mult = libfrequencies.get_freq_code(self.freq)
         result = get_period_field_arr(alias, self.asi8, base)
@@ -170,7 +164,7 @@ def __init__(self, values, freq=None, dtype=None, copy=False):
         self._dtype = PeriodDtype(freq)
 
     @classmethod
-    def _simple_new(cls, values: np.ndarray, freq=None, **kwargs):
+    def _simple_new(cls, values: np.ndarray, freq=None, **kwargs) -> "PeriodArray":
         # alias for PeriodArray.__init__
         assert isinstance(values, np.ndarray) and values.dtype == "i8"
         return cls(values, freq=freq, **kwargs)
@@ -181,7 +175,7 @@ def _from_sequence(
         scalars: Sequence[Optional[Period]],
         dtype: Optional[PeriodDtype] = None,
         copy: bool = False,
-    ) -> ABCPeriodArray:
+    ) -> "PeriodArray":
         if dtype:
             freq = dtype.freq
         else:
@@ -202,11 +196,13 @@ def _from_sequence(
         return cls(ordinals, freq=freq)
 
     @classmethod
-    def _from_sequence_of_strings(cls, strings, dtype=None, copy=False):
+    def _from_sequence_of_strings(
+        cls, strings, dtype=None, copy=False
+    ) -> "PeriodArray":
         return cls._from_sequence(strings, dtype, copy)
 
     @classmethod
-    def _from_datetime64(cls, data, freq, tz=None):
+    def _from_datetime64(cls, data, freq, tz=None) -> "PeriodArray":
         """
         Construct a PeriodArray from a datetime64 array
 
@@ -270,12 +266,12 @@ def _check_compatible_with(self, other, setitem: bool = False):
     # Data / Attributes
 
     @cache_readonly
-    def dtype(self):
+    def dtype(self) -> PeriodDtype:
         return self._dtype
 
     # error: Read-only property cannot override read-write property  [misc]
     @property  # type: ignore
-    def freq(self):
+    def freq(self) -> DateOffset:
         """
         Return the frequency object for this PeriodArray.
         """
@@ -402,7 +398,7 @@ def __arrow_array__(self, type=None):
     daysinmonth = days_in_month
 
     @property
-    def is_leap_year(self):
+    def is_leap_year(self) -> np.ndarray:
         """
         Logical indicating if the date belongs to a leap year.
         """
@@ -458,12 +454,6 @@ def to_timestamp(self, freq=None, how="start"):
         new_data = libperiod.periodarr_to_dt64arr(new_data.asi8, base)
         return DatetimeArray._from_sequence(new_data, freq="infer")
 
-    # --------------------------------------------------------------------
-    # Array-like / EA-Interface Methods
-
-    def _values_for_argsort(self):
-        return self._data
-
     # --------------------------------------------------------------------
 
     def _time_shift(self, periods, freq=None):
@@ -495,7 +485,7 @@ def _time_shift(self, periods, freq=None):
     def _box_func(self):
         return lambda x: Period._from_ordinal(ordinal=x, freq=self.freq)
 
-    def asfreq(self, freq=None, how="E"):
+    def asfreq(self, freq=None, how="E") -> "PeriodArray":
         """
         Convert the Period Array/Index to the specified frequency `freq`.
 
@@ -557,7 +547,7 @@ def asfreq(self, freq=None, how="E"):
     # ------------------------------------------------------------------
     # Rendering Methods
 
-    def _formatter(self, boxed=False):
+    def _formatter(self, boxed: bool = False):
         if boxed:
             return str
         return "'{}'".format
@@ -584,7 +574,7 @@ def _format_native_types(self, na_rep="NaT", date_format=None, **kwargs):
 
     # ------------------------------------------------------------------
 
-    def astype(self, dtype, copy=True):
+    def astype(self, dtype, copy: bool = True):
         # We handle Period[T] -> Period[U]
         # Our parent handles everything else.
         dtype = pandas_dtype(dtype)
@@ -965,8 +955,8 @@ def _get_ordinal_range(start, end, periods, freq, mult=1):
     if end is not None:
         end = Period(end, freq)
 
-    is_start_per = isinstance(start, ABCPeriod)
-    is_end_per = isinstance(end, ABCPeriod)
+    is_start_per = isinstance(start, Period)
+    is_end_per = isinstance(end, Period)
 
     if is_start_per and is_end_per and start.freq != end.freq:
         raise ValueError("start and end must have same freq")

pandas/core/ops/dispatch.py

@@ -1,10 +1,12 @@
 """
 Functions for defining unary operations.
 """
-from typing import Any, Union
+from typing import Any
 
 import numpy as np
 
+from pandas._typing import ArrayLike
+
 from pandas.core.dtypes.common import (
     is_datetime64_dtype,
     is_extension_array_dtype,
@@ -13,7 +15,7 @@
     is_scalar,
     is_timedelta64_dtype,
 )
-from pandas.core.dtypes.generic import ABCExtensionArray, ABCSeries
+from pandas.core.dtypes.generic import ABCSeries
 
 from pandas.core.construction import array
 
@@ -93,9 +95,7 @@ def should_series_dispatch(left, right, op):
     return False
 
 
-def dispatch_to_extension_op(
-    op, left: Union[ABCExtensionArray, np.ndarray], right: Any,
-):
+def dispatch_to_extension_op(op, left: ArrayLike, right: Any):
     """
     Assume that left or right is a Series backed by an ExtensionArray,
     apply the operator defined by op.

pandas/io/pytables.py

@@ -2202,7 +2202,7 @@ def __eq__(self, other: Any) -> bool:
             for a in ["name", "cname", "dtype", "pos"]
         )
 
-    def set_data(self, data: Union[np.ndarray, ABCExtensionArray]):
+    def set_data(self, data: ArrayLike):
         assert data is not None
         assert self.dtype is None
 
@@ -4959,11 +4959,11 @@ def _dtype_to_kind(dtype_str: str) -> str:
     return kind
 
 
-def _get_data_and_dtype_name(data: Union[np.ndarray, ABCExtensionArray]):
+def _get_data_and_dtype_name(data: ArrayLike):
     """
     Convert the passed data into a storable form and a dtype string.
     """
-    if is_categorical_dtype(data.dtype):
+    if isinstance(data, Categorical):
         data = data.codes
 
     # For datetime64tz we need to drop the TZ in tests TODO: why?