ci: Add mypy to ci and update types in code (#25)

Author: FBumann

.github/workflows/ci.yml

@@ -31,6 +31,9 @@ jobs:
       - name: Format check
         run: uv run ruff format --check .
 
+      - name: Type check
+        run: uv run mypy xarray_plotly
+
       - name: Test
         run: uv run pytest --cov=xarray_plotly --cov-report=xml
 

xarray_plotly/__init__.py

@@ -111,5 +111,5 @@ def xpx(data: DataArray | Dataset) -> DataArrayPlotlyAccessor | DatasetPlotlyAcc
 __version__ = version("xarray_plotly")
 
 # Register the accessors
-register_dataarray_accessor("plotly")(DataArrayPlotlyAccessor)
-register_dataset_accessor("plotly")(DatasetPlotlyAccessor)
+register_dataarray_accessor("plotly")(DataArrayPlotlyAccessor)  # type: ignore[no-untyped-call]
+register_dataset_accessor("plotly")(DatasetPlotlyAccessor)  # type: ignore[no-untyped-call]

xarray_plotly/common.py

@@ -159,7 +159,7 @@ def to_dataframe(darray: DataArray) -> pd.DataFrame:
     return df
 
 
-def _get_label_from_attrs(attrs: dict, fallback: str) -> str:
+def _get_label_from_attrs(attrs: dict[str, object], fallback: str) -> str:
     """Extract a label from xarray attributes based on current config.
 
     Args:

xarray_plotly/config.py

@@ -8,7 +8,7 @@
 
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 if TYPE_CHECKING:
     from collections.abc import Generator
@@ -166,12 +166,12 @@ def set_options(
         yield
     finally:
         # Restore old values (modify in place)
-        _options.label_use_long_name = old_values["label_use_long_name"]
-        _options.label_use_standard_name = old_values["label_use_standard_name"]
-        _options.label_include_units = old_values["label_include_units"]
-        _options.label_unit_format = old_values["label_unit_format"]
-        _options.slot_orders = old_values["slot_orders"]
-        _options.dataset_variable_position = old_values["dataset_variable_position"]
+        _options.label_use_long_name = cast("bool", old_values["label_use_long_name"])
+        _options.label_use_standard_name = cast("bool", old_values["label_use_standard_name"])
+        _options.label_include_units = cast("bool", old_values["label_include_units"])
+        _options.label_unit_format = cast("str", old_values["label_unit_format"])
+        _options.slot_orders = cast("dict[str, tuple[str, ...]]", old_values["slot_orders"])
+        _options.dataset_variable_position = cast("int", old_values["dataset_variable_position"])
 
 
 def notebook(renderer: str = "notebook") -> None:

xarray_plotly/figures.py

@@ -5,15 +5,15 @@
 from __future__ import annotations
 
 import copy
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
     import plotly.graph_objects as go
 
 
-def _iter_all_traces(fig: go.Figure) -> Iterator:
+def _iter_all_traces(fig: go.Figure) -> Iterator[Any]:
     """Iterate over all traces in a figure, including animation frames.
 
     Yields traces from fig.data first, then from each frame in fig.frames.
@@ -107,7 +107,7 @@ def _merge_frames(
     overlays: list[go.Figure],
     base_trace_count: int,
     overlay_trace_counts: list[int],
-) -> list:
+) -> list[go.Frame]:
     """Merge animation frames from base and overlay figures.
 
     Args:
@@ -360,7 +360,7 @@ def _merge_secondary_y_frames(
     base: go.Figure,
     secondary: go.Figure,
     y_mapping: dict[str, str],
-) -> list:
+) -> list[go.Frame]:
     """Merge animation frames for secondary y-axis combination.
 
     Args:
@@ -411,7 +411,9 @@ def _merge_secondary_y_frames(
     return merged_frames
 
 
-def update_traces(fig: go.Figure, selector: dict | None = None, **kwargs) -> go.Figure:
+def update_traces(
+    fig: go.Figure, selector: dict[str, Any] | None = None, **kwargs: Any
+) -> go.Figure:
     """Update traces in both base figure and all animation frames.
 
     Plotly's `update_traces()` only updates the base figure, not animation frames.

xarray_plotly/plotting.py

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
+import numpy.typing as npt
 import plotly.express as px
 
 from xarray_plotly.common import (
@@ -191,7 +192,7 @@ def bar(
     )
 
 
-def _classify_trace_sign(y_values: np.ndarray) -> str:
+def _classify_trace_sign(y_values: npt.ArrayLike) -> str:
     """Classify a trace as 'positive', 'negative', or 'mixed' based on its values."""
     y_arr = np.asarray(y_values)
     y_clean = y_arr[np.isfinite(y_arr) & (np.abs(y_arr) > 1e-9)]