diff --git a/doc/release_notes.rst b/doc/release_notes.rst index ac36b3f1..fcdc8959 100644 --- a/doc/release_notes.rst +++ b/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 `__) **Internal** diff --git a/linopy/variables.py b/linopy/variables.py index 5b9b997d..8f3edd05 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -6,7 +6,6 @@ from __future__ import annotations -import functools import logging from collections.abc import Callable, Hashable, Mapping from dataclasses import dataclass @@ -14,6 +13,9 @@ 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)