Skip to content

Commit

Permalink
Start a list of modules which require typing (#8198)
Browse files Browse the repository at this point in the history
* Start a list of modules which require typing

Notes inline. Just one module so far!
  • Loading branch information
max-sixty authored Sep 20, 2023
1 parent 1d2e5c6 commit 8c21376
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 6 deletions.
11 changes: 10 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ show_error_context = true
warn_redundant_casts = true
warn_unused_ignores = true

# Most of the numerical computing stack doesn't have type annotations yet.
# Much of the numerical computing stack doesn't have type annotations yet.
[[tool.mypy.overrides]]
ignore_missing_imports = true
module = [
Expand Down Expand Up @@ -118,6 +118,15 @@ module = [
"numpy.exceptions.*", # remove once support for `numpy<2.0` has been dropped
]

# Gradually we want to add more modules to this list, ratcheting up our total
# coverage. Once a module is here, functions require annotations in order to
# pass mypy. It would be especially useful to have tests here, because without
# annotating test functions, we don't have a great way of testing our type
# annotations — even with just `-> None` is sufficient for mypy to check them.
[[tool.mypy.overrides]]
disallow_untyped_defs = true
module = ["xarray.core.rolling_exp"]

[tool.ruff]
builtins = ["ellipsis"]
exclude = [
Expand Down
20 changes: 15 additions & 5 deletions xarray/core/rolling_exp.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,18 +9,23 @@
from xarray.core.options import _get_keep_attrs
from xarray.core.pdcompat import count_not_none
from xarray.core.pycompat import is_duck_dask_array
from xarray.core.types import T_DataWithCoords
from xarray.core.types import T_DataWithCoords, T_DuckArray


def _get_alpha(com=None, span=None, halflife=None, alpha=None):
def _get_alpha(
com: float | None = None,
span: float | None = None,
halflife: float | None = None,
alpha: float | None = None,
) -> float:
# pandas defines in terms of com (converting to alpha in the algo)
# so use its function to get a com and then convert to alpha

com = _get_center_of_mass(com, span, halflife, alpha)
return 1 / (1 + com)


def move_exp_nanmean(array, *, axis, alpha):
def move_exp_nanmean(array: T_DuckArray, *, axis: int, alpha: float) -> np.ndarray:
if is_duck_dask_array(array):
raise TypeError("rolling_exp is not currently support for dask-like arrays")
import numbagg
Expand All @@ -32,15 +37,20 @@ def move_exp_nanmean(array, *, axis, alpha):
return numbagg.move_exp_nanmean(array, axis=axis, alpha=alpha)


def move_exp_nansum(array, *, axis, alpha):
def move_exp_nansum(array: T_DuckArray, *, axis: int, alpha: float) -> np.ndarray:
if is_duck_dask_array(array):
raise TypeError("rolling_exp is not currently supported for dask-like arrays")
import numbagg

return numbagg.move_exp_nansum(array, axis=axis, alpha=alpha)


def _get_center_of_mass(comass, span, halflife, alpha):
def _get_center_of_mass(
comass: float | None,
span: float | None,
halflife: float | None,
alpha: float | None,
) -> float:
"""
Vendored from pandas.core.window.common._get_center_of_mass
Expand Down

0 comments on commit 8c21376

Please sign in to comment.