From 5d6bb592fcc60e29c66d13bcb8122c0d1dbb8ff2 Mon Sep 17 00:00:00 2001
From: Siddhartha Gandhi <siddhartha.a.gandhi@gmail.com>
Date: Wed, 7 Aug 2024 12:15:51 -0400
Subject: [PATCH] Annotate depends and accept_arguments decorators

---
 param/_utils.py  | 12 ++++++++++--
 param/depends.py | 34 ++++++++++++++++++++++++++++++++--
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/param/_utils.py b/param/_utils.py
index 4fdc71c9..1e921d9a 100644
--- a/param/_utils.py
+++ b/param/_utils.py
@@ -10,6 +10,7 @@
 import sys
 import traceback
 import warnings
+from typing import ParamSpec, TypeVar, Callable, TYPE_CHECKING, Concatenate
 
 from collections import defaultdict, OrderedDict
 from contextlib import contextmanager
@@ -18,6 +19,11 @@
 from threading import get_ident
 from collections import abc
 
+if TYPE_CHECKING:
+    _P1 = ParamSpec("P1")
+    _P2 = ParamSpec("P2")
+    _R = TypeVar("_R")
+
 DEFAULT_SIGNATURE = inspect.Signature([
     inspect.Parameter('self', inspect.Parameter.POSITIONAL_OR_KEYWORD),
     inspect.Parameter('params', inspect.Parameter.VAR_KEYWORD),
@@ -282,12 +288,14 @@ def flatten(line):
             yield element
 
 
-def accept_arguments(f):
+def accept_arguments(
+    f: Callable[Concatenate[Callable[_P1, _R], _P2], _R]
+) -> Callable[_P2, Callable[[Callable[_P1, _R]], Callable[_P1, _R]]]:
     """
     Decorator for decorators that accept arguments
     """
     @functools.wraps(f)
-    def _f(*args, **kwargs):
+    def _f(*args: _P2.args, **kwargs: _P2.kwargs) -> Callable[[Callable[_P1, _R]], Callable[_P1, _R]]:
         return lambda actual_f: f(actual_f, *args, **kwargs)
     return _f
 
diff --git a/param/depends.py b/param/depends.py
index 8b4c172d..2698033e 100644
--- a/param/depends.py
+++ b/param/depends.py
@@ -1,16 +1,46 @@
+from __future__ import annotations
+
 import inspect
 
 from collections import defaultdict
 from functools import wraps
+from typing import TYPE_CHECKING, TypeVar, Callable, Protocol, TypedDict, overload
 
 from .parameterized import (
     Parameter, Parameterized, ParameterizedMetaclass, transform_reference,
 )
 from ._utils import accept_arguments, iscoroutinefunction
 
+if TYPE_CHECKING:
+    CallableT = TypeVar("CallableT", bound=Callable)
+    Dependency = Parameter | str
+
+    class DependencyInfo(TypedDict):
+        dependencies: tuple[Dependency, ...]
+        kw: dict[str, Dependency]
+        watch: bool
+        on_init: bool
+
+    class DependsFunc(Protocol[CallableT]):
+        _dinfo: DependencyInfo
+        __call__: CallableT
+
+@overload
+def depends(
+    *dependencies: str, watch: bool = ..., on_init: bool = ...
+) -> Callable[[CallableT], DependsFunc[CallableT]]:
+    ...
+
+@overload
+def depends(
+    *dependencies: Parameter, watch: bool = ..., on_init: bool = ..., **kw: Parameter
+) -> Callable[[CallableT], DependsFunc[CallableT]]:
+    ...
 
 @accept_arguments
-def depends(func, *dependencies, watch=False, on_init=False, **kw):
+def depends(
+    func: CallableT, *dependencies: Dependency, watch: bool = False, on_init: bool = False, **kw: Parameter
+) -> Callable[[CallableT], DependsFunc[CallableT]]:
     """Annotates a function or Parameterized method to express its dependencies.
 
     The specified dependencies can be either be Parameter instances or if a
@@ -117,6 +147,6 @@ def cb(*events):
     _dinfo.update({'dependencies': dependencies,
                    'kw': kw, 'watch': watch, 'on_init': on_init})
 
-    _depends._dinfo = _dinfo
+    _depends._dinfo = _dinfo  # type: ignore[attr-defined]
 
     return _depends