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.

I believe that's where you'd use Generics and TypeVars, but I haven't used them enough to provide a spontaneous sample. — deceze, Mar 08 '21 at 08:31
The code is wrong anyway, your returns must be in `'` for example `set_width(self, width: int) -> 'M':` as you cannot type the class you are in it without it. Second of all in your calls it should be `box = M().set_width(5).set_height(10)` and lastly, once you fix that the typing works as expected, including the inheritance — bluesummers, Mar 08 '21 at 09:05
@bluesummers Thank you for pointing out the syntax errors, though the last line of code still has the incorrect types — mousetail, Mar 08 '21 at 09:09
About syntax you should add `'` aroudn the class return types as I showed in the last comment. Other than that, in my local environent, there are zero typing problems — bluesummers, Mar 08 '21 at 09:10
You don't need to, you can simply import `annotations` from `__future__` — mousetail, Mar 08 '21 at 09:10
Well, this code has no errors for me and PyCharm highllights everything as working... What is the error yo uare getting? — bluesummers, Mar 08 '21 at 09:12
I mean, it won't error, but PyCharm won't autocomplete any methods on cube since it can't determine the type of the variable since it thinks the method `set_depth()` doesn't exist — mousetail, Mar 08 '21 at 09:13
Without using `annotations` (using the `'`) works fine and autocompletes in PyCharm for me. Could be that PyCharm did not develop this feature yet — bluesummers, Mar 08 '21 at 09:17

score 1 · Answer 1 · answered Mar 08 '21 at 08:54

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).

score 0 · Accepted Answer · answered Mar 15 '21 at 14:37

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.

Python how to type anotate a method that returns self?

2 Answers2