TYP: numba stub (#44233)

Author: simonjayhawkins

ci/code_checks.sh

@@ -106,7 +106,7 @@ if [[ -z "$CHECK" || "$CHECK" == "typing" ]]; then
     mypy --version
 
     MSG='Performing static analysis using mypy' ; echo $MSG
-    mypy pandas
+    mypy
     RET=$(($RET + $?)) ; echo $MSG "DONE"
 
     # run pyright, if it is installed

doc/source/development/contributing_codebase.rst

@@ -399,7 +399,7 @@ pandas uses `mypy <http://mypy-lang.org>`_ and `pyright <https://github.com/micr
 
 .. code-block:: shell
 
-   mypy pandas
+   mypy
 
    # let pre-commit setup and run pyright
    pre-commit run --hook-stage manual --all-files pyright

pandas/core/_numba/kernels/mean_.py

@@ -14,8 +14,7 @@
 from pandas.core._numba.kernels.shared import is_monotonic_increasing
 
 
-# error: Untyped decorator makes function "add_mean" untyped
-@numba.jit(nopython=True, nogil=True, parallel=False)  # type: ignore[misc]
+@numba.jit(nopython=True, nogil=True, parallel=False)
 def add_mean(
     val: float, nobs: int, sum_x: float, neg_ct: int, compensation: float
 ) -> tuple[int, float, int, float]:
@@ -30,8 +29,7 @@ def add_mean(
     return nobs, sum_x, neg_ct, compensation
 
 
-# error: Untyped decorator makes function "remove_mean" untyped
-@numba.jit(nopython=True, nogil=True, parallel=False)  # type: ignore[misc]
+@numba.jit(nopython=True, nogil=True, parallel=False)
 def remove_mean(
     val: float, nobs: int, sum_x: float, neg_ct: int, compensation: float
 ) -> tuple[int, float, int, float]:
@@ -46,8 +44,7 @@ def remove_mean(
     return nobs, sum_x, neg_ct, compensation
 
 
-# error: Untyped decorator makes function "sliding_mean" untyped
-@numba.jit(nopython=True, nogil=True, parallel=False)  # type: ignore[misc]
+@numba.jit(nopython=True, nogil=True, parallel=False)
 def sliding_mean(
     values: np.ndarray,
     start: np.ndarray,

pandas/core/_numba/kernels/shared.py

@@ -2,9 +2,12 @@
 import numpy as np
 
 
-# error: Untyped decorator makes function "is_monotonic_increasing" untyped
-@numba.jit(  # type: ignore[misc]
-    numba.boolean(numba.int64[:]), nopython=True, nogil=True, parallel=False
+@numba.jit(
+    # error: Any? not callable
+    numba.boolean(numba.int64[:]),  # type: ignore[misc]
+    nopython=True,
+    nogil=True,
+    parallel=False,
 )
 def is_monotonic_increasing(bounds: np.ndarray) -> bool:
     """Check if int64 values are monotonically increasing."""

pandas/core/_numba/kernels/sum_.py

@@ -14,8 +14,7 @@
 from pandas.core._numba.kernels.shared import is_monotonic_increasing
 
 
-# error: Untyped decorator makes function "add_sum" untyped
-@numba.jit(nopython=True, nogil=True, parallel=False)  # type: ignore[misc]
+@numba.jit(nopython=True, nogil=True, parallel=False)
 def add_sum(
     val: float, nobs: int, sum_x: float, compensation: float
 ) -> tuple[int, float, float]:
@@ -28,8 +27,7 @@ def add_sum(
     return nobs, sum_x, compensation
 
 
-# error: Untyped decorator makes function "remove_sum" untyped
-@numba.jit(nopython=True, nogil=True, parallel=False)  # type: ignore[misc]
+@numba.jit(nopython=True, nogil=True, parallel=False)
 def remove_sum(
     val: float, nobs: int, sum_x: float, compensation: float
 ) -> tuple[int, float, float]:
@@ -42,8 +40,7 @@ def remove_sum(
     return nobs, sum_x, compensation
 
 
-# error: Untyped decorator makes function "sliding_sum" untyped
-@numba.jit(nopython=True, nogil=True, parallel=False)  # type: ignore[misc]
+@numba.jit(nopython=True, nogil=True, parallel=False)
 def sliding_sum(
     values: np.ndarray,
     start: np.ndarray,

pyproject.toml

@@ -60,6 +60,8 @@ markers = [
 
 [tool.mypy]
 # Import discovery
+mypy_path = "typings"
+files = ["pandas", "typings"]
 namespace_packages = false
 explicit_package_bases = false
 ignore_missing_imports = true

typings/numba.pyi

@@ -0,0 +1,41 @@
+from typing import (
+    Any,
+    Callable,
+    Literal,
+    overload,
+)
+
+import numba
+
+from pandas._typing import F
+
+def __getattr__(name: str) -> Any: ...  # incomplete
+@overload
+def jit(
+    signature_or_function: F = ...,
+) -> F: ...
+@overload
+def jit(
+    signature_or_function: str
+    | list[str]
+    | numba.core.types.abstract.Type
+    | list[numba.core.types.abstract.Type] = ...,
+    locals: dict = ...,  # TODO: Mapping of local variable names to Numba types
+    cache: bool = ...,
+    pipeline_class: numba.compiler.CompilerBase = ...,
+    boundscheck: bool | None = ...,
+    *,
+    nopython: bool = ...,
+    forceobj: bool = ...,
+    looplift: bool = ...,
+    error_model: Literal["python", "numpy"] = ...,
+    inline: Literal["never", "always"] | Callable = ...,
+    # TODO: If a callable is provided it will be called with the call expression
+    # node that is requesting inlining, the caller's IR and callee's IR as
+    # arguments, it is expected to return Truthy as to whether to inline.
+    target: Literal["cpu", "gpu", "npyufunc", "cuda"] = ...,  # deprecated
+    nogil: bool = ...,
+    parallel: bool = ...,
+) -> Callable[[F], F]: ...
+
+njit = jit