Automatically infer ParamSpec parameter from a method implementation?

[alanhdu]

I have a pattern that looks something like this:

import abc

class BaseClass(metaclass=abc.ABCMeta)
   @abc.abstractmethod
   def forward(self, *args, **kwargs):
      """Function to override to implement the core logic"""

   def wrapper(self, *args, **kwargs):
      """Some wrapper around `.forward`"""
      return self.forward(*args, **kwargs)

and would like wrapper to have the same type annotation as forward (this is a common pattern that is used in PyTorch and TensorFlow when defining custom modules, for instance).

I can almost get what I want using ParamSpec and TypeVar like:

P = ParamSpec("P")
T = TypeVar("T")
class BaseClass(Generic[P, T], metaclass=abc.ABCMeta)
   @abc.abstractmethod
   def forward(self, *args: P.args, **kwargs: P.kwargs) -> T:
      """Function to override to implement the core logic"""

   def wrapper(self, *args: P.args, **kwargs: P.kwargs) -> T:
      """Some wrapper around `.forward`"""
      return self.forward(*args, **kwargs)

But this is not quite what I want:

• There's some boilerplate in the subclass, where I have to explicitly define P and T again in the inheritance call (e.g. class Subclass(BaseClass[[int], str]) along with def forward(self, x: int, /) -> str
• AFAIK, there is no good way to define kwargs with ParamSpec, so kwargs would not be allowed.

Is there a good way to support this pattern?

[khoover]

If you can switch to composition instead of inheritance, I think you can make it work.

Params = ParamSpec("Params")
Return = TypeVar("Return", covariant=True)

class ForwardProtocol(Protocol, Generic[Params, Return]):
  def forward(self, *args: Params.args, **kwargs: Params.kwargs) -> Return:
    ...

class Wrapper(Generic[Params, Return]):
  def __init__(self, impl: ForwardProtocol[Params, Return]) -> None:
    self._impl = impl

  def wrapper(self, *args: Params.args, **kwargs: Params.kwargs) -> Return:
    return self._impl.forward(*args, **kwargs)