Python how to type anotate a method that returns self?

Question

Suppose I have a class that implements method chaining:

from __future__ import annotations

class M:
    def set_width(self, width: int)->M:
        self.width = width
        return self

    def set_height(self, height: int)->M:
        self.height = height
        return self

I could use it like this:

box = M().set_width(5).set_height(10)

This works, but if I have a subclass M3D:

class M3D(M):
    def set_depth(self, depth: int) -> M3D:
        self.depth = depth
        return self

Now I can’t do this:

cube = M3D().set_width(2).set_height(3).set_depth(5)

I get the following error in mypy:

_test_typeanotations.py:21: error: "M" has no attribute "set_depth"; maybe "set_width"

Because set_width() returns an M which has no method set_depth. I have seen suggestions to override set_width() and set_height() for every subclass to specify the correct types, but that would be a lot of code to write for each method. There has to be a easier way.

This is also relevant for special methods, for example __enter__ traditionally returns self, so it would be nice to have a way to specify this without needing to even mention it in subclasses.

Asked By: mousetail

||

Source

Answer 1

This is a classic problem in any language using inheritance. And it is handled differently by the languages:

in C++, you would cast the result of set_height before calling set_depth
in Java, you could either use the same casting as C++, or have the IDE to generate the bunch of overrides and only manually change the type in the overriden methods.

Python is a dynamically typed language, so there is no cast instruction. So you are left with 3 possible ways:

the brave way: override all the relevant methods to call the base method and declare the new type in the return annotation
the I don’t care way: annotation control only gives warnings. As you know that the line is fine, you can just ignore the warning
the don’t bother me way: annotations are optional in Python and annotation control can normally be suspended by special comments. Here you know that there is no problem, so you can safely suspend the type control for that instruction or that method.

The following is only my opinion.

I would avoid the don’t bother way if possible, because if you will leave warnings in your code, you would have to later control after each and every change if there is a new warning.

I would not override methods just to get rid of a warning. After all Python is a dynamically typed language that even allows duck typing. If I know that the code is correct I would avoid adding useless code (DRY and KISS principles)

SO I will just assume that comments to suspend annotation controls were invented for a reason and use them (what I call don’t bother me here).

Answered By: Serge Ballesta

Answer 2

After a lot of research and expirimentation, I have found a way that works in mypy, though Pycham still guesses the type wrong sometimes.

The trick is to make self a type var:

from __future__ import annotations

import asyncio
from typing import TypeVar

T = TypeVar('T')


class M:
    def set_width(self: T, width: int)->T:
        self.width = width
        return self

    def set_height(self: T, height: int)->T:
        self.height = height
        return self

    def copy(self)->M:
        return M().set_width(self.width).set_height(self.height)


class M3D(M):
    def set_depth(self: T, depth: int) -> T:
        self.depth = depth
        return self

box = M().set_width(5).set_height(10) # box has correct type
cube = M3D().set_width(2).set_height(3).set_depth(5) # cube has correct type
attemptToTreatBoxAsCube = M3D().copy().set_depth(4) # Mypy gets angry as expected

The last line specifically works fine in mypy but pycharm will still autocomplete set_depth sometimes even though .copy() actually returns an M even when called on a M3D.

Answered By: mousetail

Answer 3

In Python 3.11 and its later versions, you will be able to do this:

from typing import Self

class M:
    def set_width(self, width: int) -> Self:
        self.width = width
        return self

Answered By: Boštjan Mejak

Python how to type anotate a method that returns self?

Question:

Answers: