CODEWITHSUNDEEP
pythonmypypython-typing
Cesar Canassa
Cesar Canassa

Reputation: 20203

How to override a class __init__ method while keeping the types

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.

Upvotes: 3

Views: 1481

Answers (3)

PhantomAmnesiac
PhantomAmnesiac

Reputation: 41

In case someone is looking for a way to have VSCode show a signature for __init__ that includes arguments from the parent class, as well as the subclass, here is what I found:

I was able to achieve this using the dataclass module.

from dataclasses import dataclass


@dataclass
class Base:
    """Base model to be extended"""

    arg1: bool
    arg2: bool

    def __post_init__(self) -> None:
        # Called after __init__ method
        pass


@dataclass
class SubClass(Base):
    """Child model extending Base"""

    arg3: bool
    arg4: bool

    def __post_init__(self) -> None:
        # Called after __init__ method
        pass

Now, VSCode shows me a method signature with all four arguments when I hover over SubClass. I don't have to use *args, **kwargs.

My gripe with this solution is that I cannot have default arguments in Base class, if I have non-default arguments in SubClass

Upvotes: 1

Robin Gugel
Robin Gugel

Reputation: 1216

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.

Upvotes: 3

Jay M
Jay M

Reputation: 4317

The real answer depends on your use case. You can (for instance) implement the __new__() special method on your class to create an object of a given type dependent on the parameter provided to __init__().

In most cases this is rather too complicated and Python provides a generic means for this that covers a lot of use cases. Take a look at Python DataClasses

The other way might be to review the more general response to this question that I wrote some time back. "SubClassing Int in Python", specifically the example of a Modified Type class.

Upvotes: -1

Related Questions