Inherit typing from decorated function in @named_scope

src/jaxsim/api/common.py

@@ -3,7 +3,8 @@
 import dataclasses
 import enum
 import functools
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
+from typing import ParamSpec, TypeVar
 
 import jax
 import jax.numpy as jnp
@@ -20,11 +21,15 @@
     from typing_extensions import Self
 
 
-def named_scope(fn, name: str | None = None):
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+def named_scope(fn, name: str | None = None) -> Callable[_P, _R]:
     """Applies a JAX named scope to a function for improved profiling and clarity."""
 
     @functools.wraps(fn)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
         with jax.named_scope(name or fn.__name__):
             return fn(*args, **kwargs)