Type checking your classes
The next step up from functions, unsurprisingly, is classes. Fortunately type checking classes is very straightforward: Just add hints to the methods for the most part and you'll be good.
Here's an example class, that just implements a stack of numbers:
class Stack:
def __init__(self):
self._values = []
def __repr__(self):
return f'Stack{self._values!r}'
def push(self, value):
self._values.append(value)
def pop(self):
if len(self._values) == 0:
raise RuntimeError('Underflow!')
return self._values.pop()
stack = Stack()
print(stack)
stack.push(2)
stack.push(10)
print(stack)
popped_value = stack.pop()
print("Popped value:", popped_value)
print(stack)
Running the code, you can see that the stack gets values put into it and popped out of it. If you pop on an empty stack, you'll get an Underflow error.
Let's see what mypy says about it, and add type annotations as necessary:
First off, mypy complains about a lot of untyped definitions. The functions need annotation. And before you do that, think about this:
def __init__(self):
What should the type of self
be?
Since it's an object of the Stack
type, it would make sense to do:
def __init__(self: Stack) -> None:
Right? Well, kindof. The problem then arises when you try to subclass Stack
into some other type:
class SubStack(Stack):
...
Since the __init__
method is inherited, and self
is supposed to be of type
Stack
, mypy would thinkg SubStack
creates objects of Stack
type. Clearly
that won't work.
To avoid problems with inheritence, mypy decided that self
does not need to be
given a type annotation. You must leave the first argument of a method untyped.
So our annotations look something like this:
def push(self, value: int) -> None:
Here's what the class looks like after annotating the function signatures:
class Stack:
def __init__(self) -> None:
self._values = []
def __repr__(self) -> str:
return f'Stack{self._values!r}'
def push(self, value: int) -> None:
self._values.append(value)
def pop(self) -> int:
if len(self._values) == 0:
raise RuntimeError('Underflow!')
return self._values.pop()
Although looking at mypy's output, it still isn't sure about some things. Can you guess why?
It's because of the empty list instantiation inside __init__
. Mypy needs to
know what kinds of items you want to be in the list. The second error stems from
the same fact: Since mypy doesn't know what the type of self._values
is, it
can't determine what self._values.pop()
should return.
To fix that, replace the definition with this:
self._values: list[int] = []
And now we have fully annotated a class. Wasn't it really easy? That's pretty much how it goes with mypy, 90% of your work is adding types to function definitions, and mypy takes care of validating the rest.
Aside: exceptions
Notice the definition of pop
in the Stack
class:
def pop(self) -> int:
if len(self._values) == 0:
raise RuntimeError('Underflow!')
return self._values.pop()
In case of an underflow, the function raises an exception. The function doesn't
return an int
in that case. So you might think that the type signature isn't
fully correct: Sometimes we return int
, but sometimes we don't. And yet mypy
is fully convinced that the signature is perfectly good. Why is that?
Take this code for example:
def only_evens(number: int) -> int:
if number % 2 != 0:
raise ValueError("Not even!")
return number
number = 2
even_number = only_evens(number)
print(even_number, "is an even number.")
Mypy says the function is okay. The code also works, but if you change number
to be 3
, we can see that the code crashes. What's going on?
Turns out, mypy is right (and almost always, it will be). When only_evens
raises an error, the normal execution of the code stops, and the error is
propagated up to the user. Which means that if number
is odd, because of the
raised error, even_number
variable will never be created.
The important thing to note about this code snippet is this: If even_number
is
ever created, it is guaranteed to be an int
. So the return type of the
function is accurate. Whenever the function returns a value, it is 100%
guaranteed to be int
. The types are correct.