Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Start a list of modules which require typing #8198

Merged
merged 4 commits into from
Sep 20, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading