annotate __contains__

Author: jbrockmendel

pandas/core/indexes/category.py

@@ -1,4 +1,4 @@
-from typing import Any, List
+from typing import Any, Hashable, List, Optional
 import warnings
 
 import numpy as np
@@ -385,7 +385,7 @@ def _wrap_setop_result(self, other, result):
         return self._shallow_copy(result, name=name)
 
     @Appender(_index_shared_docs["contains"] % _index_doc_kwargs)
-    def __contains__(self, key) -> bool:
+    def __contains__(self, key: Optional[Hashable]) -> bool:
         # if key is a NaN, check if any NaN is in self.
         if is_scalar(key) and isna(key):
             return self.hasnans

pandas/core/indexes/datetimelike.py

@@ -2,7 +2,7 @@
 Base and utility classes for tseries type pandas objects.
 """
 import operator
-from typing import List, Optional, Set
+from typing import Hashable, List, Optional, Set
 
 import numpy as np
 
@@ -153,7 +153,7 @@ def equals(self, other) -> bool:
         return np.array_equal(self.asi8, other.asi8)
 
     @Appender(_index_shared_docs["contains"] % _index_doc_kwargs)
-    def __contains__(self, key):
+    def __contains__(self, key: Optional[Hashable]):
         hash(key)
         try:
             res = self.get_loc(key)

pandas/core/indexes/interval.py

@@ -1,7 +1,7 @@
 """ define the IntervalIndex """
 from operator import le, lt
 import textwrap
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Hashable, Optional, Tuple, Union
 
 import numpy as np
 
@@ -372,7 +372,7 @@ def _engine(self):
         right = self._maybe_convert_i8(self.right)
         return IntervalTree(left, right, closed=self.closed)
 
-    def __contains__(self, key) -> bool:
+    def __contains__(self, key: Optional[Hashable]) -> bool:
         """
         return a boolean if this key is IN the index
         We *only* accept an Interval

pandas/core/indexes/numeric.py

@@ -1,3 +1,5 @@
+from typing import Hashable, Optional
+
 import numpy as np
 
 from pandas._libs import index as libindex, lib
@@ -225,7 +227,7 @@ class IntegerIndex(NumericIndex):
     This is an abstract class for Int64Index, UInt64Index.
     """
 
-    def __contains__(self, key) -> bool:
+    def __contains__(self, key: Optional[Hashable]) -> bool:
         """
         Check if key is a float and has a decimal. If it has, return False.
         """
@@ -473,7 +475,7 @@ def equals(self, other) -> bool:
         except (TypeError, ValueError):
             return False
 
-    def __contains__(self, other) -> bool:
+    def __contains__(self, other: Optional[Hashable]) -> bool:
         hash(other)
         if super().__contains__(other):
             return True

pandas/core/indexes/period.py

@@ -1,4 +1,5 @@
 from datetime import datetime, timedelta
+from typing import Hashable, Optional
 import weakref
 
 import numpy as np
@@ -370,7 +371,7 @@ def _engine(self):
         return self._engine_type(period, len(self))
 
     @Appender(_index_shared_docs["contains"])
-    def __contains__(self, key) -> bool:
+    def __contains__(self, key: Optional[Hashable]) -> bool:
         if isinstance(key, Period):
             if key.freq != self.freq:
                 return False