How to override a class __init__ method while keeping the types

Question

What's the proper way to extend a class __init__ method while keeping the type annotations intact?

Take this example class:

class Base:
    def __init__(self, *, a: str):
        pass

I would like to subclass Base and add a new parameter b to the __init__ method:

from typing import Any

class Sub(Base):
    def __init__(self, *args: Any, b: str, **kwargs: Any):
        super().__init__(*args, **kwargs)

The problem with that approach is that now Sub basically accepts anything. For example, mypy will happily accept the following:

Sub(a="", b="", invalid=1). # throws __init__() got an unexpected keyword argument 'invalid'

I also don't want to redefine a in Sub, since Base might be an external library that I don't fully control.

Answer 1

There is a solution for the question of "adding parameters to a method signature" - but it's not pretty... Using ParamSpecs and Concatenate you can essentially capture the parameters of your Base init and extend them.

Concatenate only enables adding new positional arguments though. The reasoning for that is stated in the PEP introducing the ParamSpec. In short, when adding a keyword-parameter we would run into problems if that keyword-parameter is already used by the function we're extending.

Check out this code. It's quite advanced but that way you can keep the type annotation of your Base class init without rewriting them.

from typing import Callable, Type, TypeVar, overload
from typing_extensions import ParamSpec, Concatenate


P = ParamSpec("P")

TSelf = TypeVar("TSelf")
TReturn = TypeVar("TReturn")
T0 = TypeVar("T0")
T1 = TypeVar("T1")
T2 = TypeVar("T2")

@overload
def add_args_to_signature(
   to_signature: Callable[Concatenate[TSelf, P], TReturn],
   new_arg_type: Type[T0]
) -> Callable[
   [Callable[..., TReturn]], 
   Callable[Concatenate[TSelf, T0, P], TReturn]
]:
   pass

@overload
def add_args_to_signature(
   to_signature: Callable[Concatenate[TSelf, P], TReturn],
   new_arg_type0: Type[T0],
   new_arg_type1: Type[T1],
) -> Callable[
   [Callable[..., TReturn]], 
   Callable[Concatenate[TSelf, T0, T1, P], TReturn]
]:
   pass


@overload
def add_args_to_signature(
   to_signature: Callable[Concatenate[TSelf, P], TReturn],
   new_arg_type0: Type[T0],
   new_arg_type1: Type[T1],
   new_arg_type2: Type[T2]
) -> Callable[
   [Callable[..., TReturn]], 
   Callable[Concatenate[TSelf, T0, T1, P], TReturn]
]:
   pass

# repeat if you want to enable adding more parameters...

def add_args_to_signature(
   *_, **__
):
   return lambda f: f

class Base:
   def __init__(self, some_arg: float, *, some_kwarg: int):
      pass


class Sub(Base):

   # Note: you'll lose the name of your new args in your code editor. 
   @add_args_to_signature(Base.__init__, str)
   def __init__(self, you_can_only_add_positional_args: str, /, *args, **kwargs):
      super().__init__(*args, **kwargs)
   

Sub("hello", 3.5, some_kwarg=5)

VS-Code gives the following type hints for Sub: Sub(str, some_arg: float, *, some_kwarg: int)

I don't know though if mypy works with ParamSpec and Concatenate...

Due to a bug in VS-Code the position of the parameters aren't correctly matched though (the parameter being set is off by one).

Note that this is quite an advanced use of the typing module. You can ping me in the comments if you need additional explanations.