TYP: annotations in core.indexes (pandas-dev#29656)

Author: jbrockmendel

pandas/core/indexes/api.py

@@ -1,4 +1,5 @@
 import textwrap
+from typing import List, Set
 import warnings
 
 from pandas._libs import NaT, lib
@@ -64,17 +65,18 @@
 ]
 
 
-def get_objs_combined_axis(objs, intersect=False, axis=0, sort=True):
+def get_objs_combined_axis(
+    objs, intersect: bool = False, axis=0, sort: bool = True
+) -> Index:
     """
     Extract combined index: return intersection or union (depending on the
     value of "intersect") of indexes on given axis, or None if all objects
     lack indexes (e.g. they are numpy arrays).
 
     Parameters
     ----------
-    objs : list of objects
-        Each object will only be considered if it has a _get_axis
-        attribute.
+    objs : list
+        Series or DataFrame objects, may be mix of the two.
     intersect : bool, default False
         If True, calculate the intersection between indexes. Otherwise,
         calculate the union.
@@ -87,26 +89,27 @@ def get_objs_combined_axis(objs, intersect=False, axis=0, sort=True):
     -------
     Index
     """
-    obs_idxes = [obj._get_axis(axis) for obj in objs if hasattr(obj, "_get_axis")]
-    if obs_idxes:
-        return _get_combined_index(obs_idxes, intersect=intersect, sort=sort)
+    obs_idxes = [obj._get_axis(axis) for obj in objs]
+    return _get_combined_index(obs_idxes, intersect=intersect, sort=sort)
 
 
-def _get_distinct_objs(objs):
+def _get_distinct_objs(objs: List[Index]) -> List[Index]:
     """
     Return a list with distinct elements of "objs" (different ids).
     Preserves order.
     """
-    ids = set()
+    ids: Set[int] = set()
     res = []
     for obj in objs:
-        if not id(obj) in ids:
+        if id(obj) not in ids:
             ids.add(id(obj))
             res.append(obj)
     return res
 
 
-def _get_combined_index(indexes, intersect=False, sort=False):
+def _get_combined_index(
+    indexes: List[Index], intersect: bool = False, sort: bool = False
+) -> Index:
     """
     Return the union or intersection of indexes.
 
@@ -147,7 +150,7 @@ def _get_combined_index(indexes, intersect=False, sort=False):
     return index
 
 
-def union_indexes(indexes, sort=True):
+def union_indexes(indexes, sort=True) -> Index:
     """
     Return the union of indexes.
 
@@ -173,7 +176,7 @@ def union_indexes(indexes, sort=True):
 
     indexes, kind = _sanitize_and_check(indexes)
 
-    def _unique_indices(inds):
+    def _unique_indices(inds) -> Index:
         """
         Convert indexes to lists and concatenate them, removing duplicates.
 

pandas/core/indexes/base.py

@@ -1650,7 +1650,7 @@ def _get_grouper_for_level(self, mapper, level=None):
     # Introspection Methods
 
     @property
-    def is_monotonic(self):
+    def is_monotonic(self) -> bool:
         """
         Alias for is_monotonic_increasing.
         """
@@ -1691,7 +1691,7 @@ def is_monotonic_decreasing(self) -> bool:
         return self._engine.is_monotonic_decreasing
 
     @property
-    def _is_strictly_monotonic_increasing(self):
+    def _is_strictly_monotonic_increasing(self) -> bool:
         """
         Return if the index is strictly monotonic increasing
         (only increasing) values.
@@ -1708,7 +1708,7 @@ def _is_strictly_monotonic_increasing(self):
         return self.is_unique and self.is_monotonic_increasing
 
     @property
-    def _is_strictly_monotonic_decreasing(self):
+    def _is_strictly_monotonic_decreasing(self) -> bool:
         """
         Return if the index is strictly monotonic decreasing
         (only decreasing) values.
@@ -1725,7 +1725,7 @@ def _is_strictly_monotonic_decreasing(self):
         return self.is_unique and self.is_monotonic_decreasing
 
     @cache_readonly
-    def is_unique(self):
+    def is_unique(self) -> bool:
         """
         Return if the index has unique values.
         """
@@ -1735,22 +1735,22 @@ def is_unique(self):
     def has_duplicates(self) -> bool:
         return not self.is_unique
 
-    def is_boolean(self):
+    def is_boolean(self) -> bool:
         return self.inferred_type in ["boolean"]
 
-    def is_integer(self):
+    def is_integer(self) -> bool:
         return self.inferred_type in ["integer"]
 
-    def is_floating(self):
+    def is_floating(self) -> bool:
         return self.inferred_type in ["floating", "mixed-integer-float", "integer-na"]
 
-    def is_numeric(self):
+    def is_numeric(self) -> bool:
         return self.inferred_type in ["integer", "floating"]
 
-    def is_object(self):
+    def is_object(self) -> bool:
         return is_object_dtype(self.dtype)
 
-    def is_categorical(self):
+    def is_categorical(self) -> bool:
         """
         Check if the Index holds categorical data.
 
@@ -1786,10 +1786,10 @@ def is_categorical(self):
         """
         return self.inferred_type in ["categorical"]
 
-    def is_interval(self):
+    def is_interval(self) -> bool:
         return self.inferred_type in ["interval"]
 
-    def is_mixed(self):
+    def is_mixed(self) -> bool:
         return self.inferred_type in ["mixed"]
 
     def holds_integer(self):
@@ -1868,7 +1868,7 @@ def _isnan(self):
     @cache_readonly
     def _nan_idxs(self):
         if self._can_hold_na:
-            w, = self._isnan.nonzero()
+            w = self._isnan.nonzero()[0]
             return w
         else:
             return np.array([], dtype=np.int64)
@@ -4086,13 +4086,13 @@ def _assert_can_do_op(self, value):
             msg = "'value' must be a scalar, passed: {0}"
             raise TypeError(msg.format(type(value).__name__))
 
-    def _is_memory_usage_qualified(self):
+    def _is_memory_usage_qualified(self) -> bool:
         """
         Return a boolean if we need a qualified .info display.
         """
         return self.is_object()
 
-    def is_type_compatible(self, kind):
+    def is_type_compatible(self, kind) -> bool:
         """
         Whether the index type is compatible with the provided type.
         """
@@ -4131,14 +4131,14 @@ def is_type_compatible(self, kind):
         """
 
     @Appender(_index_shared_docs["contains"] % _index_doc_kwargs)
-    def __contains__(self, key):
+    def __contains__(self, key) -> bool:
         hash(key)
         try:
             return key in self._engine
         except (OverflowError, TypeError, ValueError):
             return False
 
-    def contains(self, key):
+    def contains(self, key) -> bool:
         """
         Return a boolean indicating whether the provided key is in the index.
 
@@ -4199,7 +4199,7 @@ def __getitem__(self, key):
         else:
             return result
 
-    def _can_hold_identifiers_and_holds_name(self, name):
+    def _can_hold_identifiers_and_holds_name(self, name) -> bool:
         """
         Faster check for ``name in self`` when we know `name` is a Python
         identifier (e.g. in NDFrame.__getattr__, which hits this to support
@@ -4290,7 +4290,7 @@ def putmask(self, mask, value):
             # coerces to object
             return self.astype(object).putmask(mask, value)
 
-    def equals(self, other):
+    def equals(self, other) -> bool:
         """
         Determine if two Index objects contain the same elements.
 
@@ -4314,7 +4314,7 @@ def equals(self, other):
             com.values_from_object(self), com.values_from_object(other)
         )
 
-    def identical(self, other):
+    def identical(self, other) -> bool:
         """
         Similar to equals, but check that other comparable attributes are
         also equal.

pandas/core/indexes/category.py

@@ -276,7 +276,7 @@ def _shallow_copy(self, values=None, dtype=None, **kwargs):
             dtype = self.dtype
         return super()._shallow_copy(values=values, dtype=dtype, **kwargs)
 
-    def _is_dtype_compat(self, other):
+    def _is_dtype_compat(self, other) -> bool:
         """
         *this is an internal non-public method*
 
@@ -407,7 +407,7 @@ def _reverse_indexer(self):
         return self._data._reverse_indexer()
 
     @Appender(_index_shared_docs["contains"] % _index_doc_kwargs)
-    def __contains__(self, key):
+    def __contains__(self, key) -> bool:
         # if key is a NaN, check if any NaN is in self.
         if is_scalar(key) and isna(key):
             return self.hasnans
@@ -455,7 +455,7 @@ def _engine(self):
 
     # introspection
     @cache_readonly
-    def is_unique(self):
+    def is_unique(self) -> bool:
         return self._engine.is_unique
 
     @property

pandas/core/indexes/datetimelike.py

@@ -148,7 +148,7 @@ def wrapper(self, other):
         return wrapper
 
     @property
-    def _ndarray_values(self):
+    def _ndarray_values(self) -> np.ndarray:
         return self._data._ndarray_values
 
     # ------------------------------------------------------------------------

pandas/core/indexes/datetimes.py

@@ -410,7 +410,7 @@ def tz(self, value):
     tzinfo = tz
 
     @cache_readonly
-    def _is_dates_only(self):
+    def _is_dates_only(self) -> bool:
         """Return a boolean if we are only dates (and don't have a timezone)"""
         from pandas.io.formats.format import _is_dates_only
 
@@ -1237,7 +1237,7 @@ def searchsorted(self, value, side="left", sorter=None):
 
         return self.values.searchsorted(value, side=side)
 
-    def is_type_compatible(self, typ):
+    def is_type_compatible(self, typ) -> bool:
         return typ == self.inferred_type or typ == "datetime"
 
     @property

pandas/core/indexes/interval.py

@@ -343,7 +343,7 @@ def _engine(self):
         right = self._maybe_convert_i8(self.right)
         return IntervalTree(left, right, closed=self.closed)
 
-    def __contains__(self, key):
+    def __contains__(self, key) -> bool:
         """
         return a boolean if this key is IN the index
         We *only* accept an Interval
@@ -483,7 +483,7 @@ def _values(self):
         return self._data
 
     @cache_readonly
-    def _ndarray_values(self):
+    def _ndarray_values(self) -> np.ndarray:
         return np.array(self._data)
 
     def __array__(self, result=None):
@@ -529,7 +529,7 @@ def inferred_type(self) -> str:
         return "interval"
 
     @Appender(Index.memory_usage.__doc__)
-    def memory_usage(self, deep=False):
+    def memory_usage(self, deep: bool = False) -> int:
         # we don't use an explicit engine
         # so return the bytes here
         return self.left.memory_usage(deep=deep) + self.right.memory_usage(deep=deep)
@@ -542,15 +542,15 @@ def mid(self):
         return self._data.mid
 
     @cache_readonly
-    def is_monotonic(self):
+    def is_monotonic(self) -> bool:
         """
         Return True if the IntervalIndex is monotonic increasing (only equal or
         increasing values), else False
         """
         return self.is_monotonic_increasing
 
     @cache_readonly
-    def is_monotonic_increasing(self):
+    def is_monotonic_increasing(self) -> bool:
         """
         Return True if the IntervalIndex is monotonic increasing (only equal or
         increasing values), else False
@@ -1213,7 +1213,7 @@ def _format_space(self):
     def argsort(self, *args, **kwargs):
         return np.lexsort((self.right, self.left))
 
-    def equals(self, other):
+    def equals(self, other) -> bool:
         """
         Determines if two IntervalIndex objects contain the same elements
         """
@@ -1374,7 +1374,7 @@ def is_all_dates(self) -> bool:
 IntervalIndex._add_logical_methods_disabled()
 
 
-def _is_valid_endpoint(endpoint):
+def _is_valid_endpoint(endpoint) -> bool:
     """helper for interval_range to check if start/end are valid types"""
     return any(
         [
@@ -1386,7 +1386,7 @@ def _is_valid_endpoint(endpoint):
     )
 
 
-def _is_type_compatible(a, b):
+def _is_type_compatible(a, b) -> bool:
     """helper for interval_range to check type compat of start/end/freq"""
     is_ts_compat = lambda x: isinstance(x, (Timestamp, DateOffset))
     is_td_compat = lambda x: isinstance(x, (Timedelta, DateOffset))