Skip to content

Commit

Permalink
TYP: numba stub (#44233)
Browse files Browse the repository at this point in the history
  • Loading branch information
simonjayhawkins authored Oct 31, 2021
1 parent 63faec3 commit 34d0498
Show file tree
Hide file tree
Showing 7 changed files with 57 additions and 17 deletions.
2 changes: 1 addition & 1 deletion ci/code_checks.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion doc/source/development/contributing_codebase.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 3 additions & 6 deletions pandas/core/_numba/kernels/mean_.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand All @@ -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]:
Expand All @@ -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,
Expand Down
9 changes: 6 additions & 3 deletions pandas/core/_numba/kernels/shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
9 changes: 3 additions & 6 deletions pandas/core/_numba/kernels/sum_.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand All @@ -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]:
Expand All @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
41 changes: 41 additions & 0 deletions typings/numba.pyi
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 34d0498

Please sign in to comment.