TYP: annotations

Author: jbrockmendel

pandas/_libs/hashtable_class_helper.pxi.in

@@ -791,7 +791,8 @@ cdef class StringHashTable(HashTable):
             raise KeyError(key)
 
     @cython.boundscheck(False)
-    def get_indexer(self, ndarray[object] values):
+    def get_indexer(self, ndarray[object] values) -> ndarray:
+        # -> np.ndarray[np.intp]
         cdef:
             Py_ssize_t i, n = len(values)
             ndarray[intp_t] labels = np.empty(n, dtype=np.intp)

pandas/_libs/lib.pyi

@@ -4,6 +4,7 @@
 from typing import (
     Any,
     Callable,
+    Generator,
 )
 
 import numpy as np
@@ -52,8 +53,7 @@ def is_bool_array(values: np.ndarray, skipna: bool = False): ...
 
 def fast_multiget(mapping: dict, keys: np.ndarray, default=np.nan) -> ArrayLike: ...
 
-# TODO: gen: Generator?
-def fast_unique_multiple_list_gen(gen: object, sort: bool = True) -> list: ...
+def fast_unique_multiple_list_gen(gen: Generator, sort: bool = True) -> list: ...
 def fast_unique_multiple_list(lists: list, sort: bool = True) -> list: ...
 def fast_unique_multiple(arrays: list, sort: bool = True) -> list: ...
 
@@ -90,10 +90,9 @@ def infer_datetimelike_array(
     arr: np.ndarray  # np.ndarray[object]
 ) -> str: ...
 
-# TODO: new_dtype -> np.dtype?
 def astype_intsafe(
     arr: np.ndarray,  # np.ndarray[object]
-    new_dtype,
+    new_dtype: np.dtype,
 ) -> np.ndarray: ...
 
 def fast_zip(ndarrays: list) -> np.ndarray: ...  # np.ndarray[object]
@@ -134,15 +133,13 @@ def memory_usage_of_objects(
 ) -> int: ...  # np.int64
 
 
-# TODO: f: Callable?
-# TODO: dtype -> DtypeObj?
 def map_infer_mask(
     arr: np.ndarray,
     f: Callable[[Any], Any],
     mask: np.ndarray,  # const uint8_t[:]
     convert: bool = ...,
     na_value: Any = ...,
-    dtype: Any = ...,
+    dtype: np.dtype = ...,
 ) -> ArrayLike: ...
 
 def indices_fast(

pandas/_libs/lib.pyx

@@ -661,7 +661,8 @@ cpdef ndarray[object] ensure_string_array(
         bint copy=True,
         bint skipna=True,
 ):
-    """Returns a new numpy array with object dtype and only strings and na values.
+    """
+    Returns a new numpy array with object dtype and only strings and na values.
 
     Parameters
     ----------
@@ -679,7 +680,7 @@ cpdef ndarray[object] ensure_string_array(
 
     Returns
     -------
-    ndarray
+    np.ndarray[object]
         An array with the input array's elements casted to str or nan-like.
     """
     cdef:
@@ -2452,7 +2453,8 @@ no_default = NoDefault.no_default  # Sentinel indicating the default value.
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def map_infer_mask(ndarray arr, object f, const uint8_t[:] mask, bint convert=True,
-                   object na_value=no_default, object dtype=object) -> "ArrayLike":
+                   object na_value=no_default, cnp.dtype dtype=np.dtype(object)
+                   ) -> "ArrayLike":
     """
     Substitute for np.vectorize with pandas-friendly dtype inference.
 
@@ -2472,7 +2474,7 @@ def map_infer_mask(ndarray arr, object f, const uint8_t[:] mask, bint convert=Tr
 
     Returns
     -------
-    ndarray
+    np.ndarray or ExtensionArray
     """
     cdef:
         Py_ssize_t i, n