feat(typing): Improve type hints for varwrap

doc/release_notes.rst

@@ -127,6 +127,7 @@ Version 1.0.0
 **Features**
 
 * :meth:`linopy.Model.restrict` cuts a subproblem out of a model, e.g. a time slice for rolling-horizon optimization. ``m.restrict(snapshot=window, fix=...)`` returns a :class:`~linopy.restrict.WindowModel`; terms on variables outside the window move to the right-hand side at their ``fix`` values. ``window.advance(...)`` moves the window and ``window.solve()`` updates the solver model in place where possible, without a warm start unless ``warm_start=True``. ``solution`` and ``dual`` are dictionaries of ``DataArray``; check ``window.status`` after solving. See the new :doc:`restrict` guide. (https://github.com/PyPSA/linopy/issues/768)
+* The ``Variable`` methods that wrap an ``xarray.Dataset`` method, such as ``sel``, ``isel``, ``shift``, ``rename`` and ``stack``, now carry the wrapped method's parameters and return ``Variable`` for type checkers and IDEs, instead of an untyped ``Callable``. (`#1029 <https://github.com/PyPSA/linopy/pull/1029>`__)
 
 **Internal**
 

linopy/variables.py

@@ -6,14 +6,16 @@
 
 from __future__ import annotations
 
-import functools
 import logging
 from collections.abc import Callable, Hashable, Mapping
 from dataclasses import dataclass
 from types import NotImplementedType
 from typing import (
     TYPE_CHECKING,
     Any,
+    Concatenate,
+    ParamSpec,
+    Protocol,
     cast,
     overload,
 )
@@ -99,20 +101,24 @@
 
 FILL_VALUE = {"labels": -1, "lower": np.nan, "upper": np.nan, "scaling": 1.0}
 
+P = ParamSpec("P")
+
+
+class DatasetMethod(Protocol[P]):
+    def __call__(self, __ds: Dataset, /, *args: P.args, **kwargs: P.kwargs) -> Any: ...
+
 
 def varwrap(
-    method: Callable, *default_args: Any, **new_default_kwargs: Any
-) -> Callable:
-    @functools.wraps(method)
-    def _varwrap(var: Variable, *args: Any, **kwargs: Any) -> Variable:
+    method: DatasetMethod[P], **new_default_kwargs: Any
+) -> Callable[Concatenate[Variable, P], Variable]:
+
+    def _varwrap(var: Variable, /, *args: P.args, **kwargs: P.kwargs) -> Variable:
         for k, v in new_default_kwargs.items():
             kwargs.setdefault(k, v)
-        return var.__class__(
-            method(var.data, *default_args, *args, **kwargs), var.model, var.name
-        )
+        return var.__class__(method(var.data, *args, **kwargs), var.model, var.name)
 
     _varwrap.__doc__ = (
-        f"Wrapper for the xarray {method.__qualname__} function for linopy.Variable"
+        f"Wrapper for the xarray {method.__qualname__} function for linopy.Variable"  # type: ignore
     )
     if new_default_kwargs:
         _varwrap.__doc__ += f" with default arguments: {new_default_kwargs}"
@@ -1474,25 +1480,25 @@ def equals(self, other: Variable) -> bool:
 
     assign_coords = varwrap(assign_flat_coords)
 
-    assign = varwrap(assign_fields)
+    assign = varwrap(assign_fields)  # type: ignore[arg-type]
 
     broadcast_like = varwrap(Dataset.broadcast_like)
 
-    compute = varwrap(Dataset.compute)
+    compute = varwrap(Dataset.compute)  # type: ignore[arg-type]
 
     # drop = varwrap(Dataset.drop)
 
-    drop_sel = varwrap(Dataset.drop_sel)
+    drop_sel = varwrap(Dataset.drop_sel)  # type: ignore[arg-type]
 
-    drop_isel = varwrap(Dataset.drop_isel)
+    drop_isel = varwrap(Dataset.drop_isel)  # type: ignore[arg-type]
 
-    expand_dims = varwrap(Dataset.expand_dims)
+    expand_dims = varwrap(Dataset.expand_dims)  # type: ignore[arg-type]
 
-    sel = varwrap(Dataset.sel)
+    sel = varwrap(Dataset.sel)  # type: ignore[arg-type]
 
-    isel = varwrap(Dataset.isel)
+    isel = varwrap(Dataset.isel)  # type: ignore[arg-type]
 
-    shift = varwrap(Dataset.shift, fill_value=_fill_value)
+    shift = varwrap(Dataset.shift, fill_value=_fill_value)  # type: ignore[arg-type]
 
     def reindex(
         self,
@@ -1536,15 +1542,15 @@ def reindex_like(
             self.name,
         )
 
-    swap_dims = varwrap(Dataset.swap_dims)
+    swap_dims = varwrap(Dataset.swap_dims)  # type: ignore[arg-type]
 
-    set_index = varwrap(Dataset.set_index)
+    set_index = varwrap(Dataset.set_index)  # type: ignore[arg-type]
 
-    rename = varwrap(Dataset.rename)
+    rename = varwrap(Dataset.rename)  # type: ignore[arg-type]
 
-    roll = varwrap(Dataset.roll)
+    roll = varwrap(Dataset.roll)  # type: ignore[arg-type]
 
-    stack = varwrap(Dataset.stack)
+    stack = varwrap(Dataset.stack)  # type: ignore[arg-type]
 
     unstack = varwrap(Dataset.unstack, fill_value=_fill_value)