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.

Answer 1

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.