Michael van Gerwen Michael van Gerwen - 1 month ago 6
Python Question

How do I specify that the return type of a method is the same as the class itself in python?

I have the following code in python 3:

class Position:

def __init__(self, x: int, y: int):
self.x = x
self.y = y

def __add__(self, other: Position) -> Position:
return Position(self.x + other.x, self.y + other.y)


But my editor (PyCharm) says that the reference Position can not be resolved (in the _add__ method). How should I specify that I expect the return type to be of type
Position
?

Edit: I think this is actually a PyCharm issue. It actually uses the information in its warnings, and code completion



But correct my if I'm wrong, and need to use some other syntax.

Answer

If you try to run this code you will get:

NameError: name 'Position' is not defined

This is because Position must be defined before you can use it on an annotation. I will go through the workarounds suggested for similar questions, if you are in a hurry "The Blessed Way"™ is the #2.

1. Define a dummy Position

Before the class definition, place a dummy definition:

class Position(object):
    pass


class Position(object):
    ...

This will get rid of the NameError and may even look OK:

>>> Position.__add__.__annotations__
{'other': __main__.Position, 'return': __main__.Position}

But is it?

>>> for k, v in Position.__add__.__annotations__.items():
...     print(k, 'is Position:', v is Position)                                                                                                                                                                                                                  
return is Position: False
other is Position: False

2. Use a string

Just use a string instead of the class itself:

...
def __add__(self, other: 'Position') -> 'Position':
   ...

If you use the Django framework this may seems familiar, as Django models use strings for forward references (foreign key definitions where the foreign model is self or is not declared yet).

Looks like this is the recommended approach according to the docs:

Forward references

When a type hint contains names that have not been defined yet, that definition may be expressed as a string literal, to be resolved later.

A situation where this occurs commonly is the definition of a container class, where the class being defined occurs in the signature of some of the methods. For example, the following code (the start of a simple binary tree implementation) does not work:

class Tree:
    def __init__(self, left: Tree, right: Tree):
    self.left = left
    self.right = right

To address this, we write:

class Tree:
    def __init__(self, left: 'Tree', right: 'Tree'):
        self.left = left
        self.right = right

The string literal should contain a valid Python expression (i.e., compile(lit, '', 'eval') should be a valid code object) and it should evaluate without errors once the module has been fully loaded. The local and global namespace in which it is evaluated should be the same namespaces in which default arguments to the same function would be evaluated.

3. Monkey-patch in order to add the annotations:

You may use this if you are using a decorator that enforces contracts:

...
    def __add__(self, other):
        return self.__class__(self.x + other.x, self.y + other.y)

Position.__add__.__annotations__['return'] = Position
Position.__add__.__annotations__['other'] = Position

At least it seems right:

>>> for k, v in Position.__add__.__annotations__.items():
...     print(k, 'is Position:', v is Position)                                                                                                                                                                                                                  
return is Position: True
other is Position: True

But thinking about it, if you are smart enough to write a contract-enforcing decorator, you probably can teach it to safely eval the annotation if it is of type buitins.str instead of builtins.type.

Conclusion

The blessed way is the #2.