Working with classes
The previous chapter explained how type hints work in general, and how you can use them to describe the contracts of your functions. This chapter is all about classes.
The basics of annotating classes are not that hard. Let's start off with an example:
class Cat:
def __init__(self, name, lives=9):
self._name = name
self._lives = lives
self._owner = None
self._kittens = []
@property
def lives(self):
return self._lives
def set_owner(self, name):
self._owner = name
def __repr__(self):
return f"<Cat {self._name!r} lives={self._lives!r} _owner={self._owner!r}>"
class Cat:
def __init__(self, name: str, lives: int = 9) -> None:
self._name = name
self._lives = lives
self._owner: str | None = None
self._kittens: list[Cat] = []
@property
def lives(self) -> int:
return self._lives
def set_owner(self, name: str) -> None:
self._owner = name
def __repr__(self) -> str:
return f"<Cat {self._name!r} lives={self._lives!r} _owner={self._owner!r}>"
As you can see:
- Methods are annotated almost exactly the same as functions
selfis not annotated in any way- Special methods like
__init__,__repr__,__len__and so on should be fully annotated as well - Attribute types are usually inferred from the
__init__method, but sometimes it's not possible. If you have an attribute that starts off asNoneor an empty collection, you should annotate it explicitly
Editor integration example
Adding kittens
Let's add some methods to work with the kittens of our cat:
def add_kitten(self, kitten: Cat) -> None:
self._kittens.append(kitten)
def adopt_kittens(self, kittens: list[Cat]) -> None:
self._kittens.extend(kittens)
def kittens(self) -> list[Cat]:
return self._kittens
Oh no, we got some red squiggles! And the code fails at runtime.
Traceback (most recent call last):
File "/home/kitty/erp-system/main.py", line 1, in <module>
class Cat:
File "/home/kitty/erp-system/main.py", line 18, in Cat
def add_kitten(self, kitten: Cat) -> None:
^^^
NameError: name 'Cat' is not defined
We said in the previous article that annotations have no effect at runtime. That's not quite true. Functions (and classes) do store their annotation as metadata, just like they store their name, default arguments and even their source code for inspection purposes.
>>> def add(x: int, y: int) -> int:
... return x + y
...
>>> add.__name__
'add'
>>> add.__annotations__
{'x': <class 'int'>, 'y': <class 'int'>, 'return': <class 'int'>}
This metadata doesn't perform any validation and doesn't make the function any faster.
However, libraries like pydantic and
typeguard can make
use of this metadata for their own purposes.
Our annotation doesn't work because we haven't created the Cat class yet, but we're
trying to reference it. It's like trying to run x = x + 1 with no x variable defined beforehand.
This can be worked around by putting the annotation in quotes:
def add_kitten(self, kitten: "Cat") -> None:
self._kittens.append(kitten)
def adopt_kittens(self, kittens: "list[Cat]") -> None:
self._kittens.extend(kittens)
def kittens(self) -> "list[Cat]":
return self._kittens
Alternative method
You may have seen this in other places:
from __future__ import annotations
class Cat:
...
def add_kitten(self, kitten: Cat) -> None:
self._kittens.append(kitten)
This is the PEP 563 way of using "deferred annotations". It automatically turns all function and class annotations into quotes. It does look cleaner and requires less work if you're rearranging classes, however, this trick is superseded by PEP 649 and will be deprecated/removed in the future.
You can use from __future__ import annotations instead, but be aware that you will have to slightly
adjust your code in future versions of Python.
Think of the kittens
The three methods we just added may look fine, but they have a few problems that can manifest when our class will be used in practice.
-
Adopting several kittens
The
adopt_kittens()method says that you need a list of kittens. That's a bit too restrictive. If I have a tuple of kittens or an iterator of kittens, should they be left cold and hungry? Or should the caller dolist(kittens)which takes extra time and memory?Moreover, it's legal for the method to modify the list provided in the argument. But that doesn't fit the intention of
adopt_kittens(). -
Getting all the kittens
The
kittens()method has similar problems, but in a more subtle way.First, it reveals our internal implementation of how we store kittens. It will be hard to switch to a
dequein the future.Secondly, if someone obtains our list of kittens, the type hints are telling them they can treat it as any old list. This can be devastating:
Iterable
Let's see how the extend() method of list objects handles the first problem. After all, it works with
lists, tuples, ranges and other things.
list.extend works with any iterable object. An
iterable is something you can put in a for loop to go over ("iterate over") its items. Iterable and
similar items like Sequence, Iterator, Mapping are available in the standard
collections.abc
module. Let's see it in action:
from collections.abc import Iterable
def handle_numbers(numbers: Iterable[int]) -> None:
numbers.append(42) # error: `Iterable[int]` has no method `append`
for n in numbers: # ok, `n` is an `int`
print(n + 1) # ok
print(n + "!") # error: can't add a string to an integer
handle_numbers([1, 2, 3]) # ok
handle_numbers((1, 2, 3)) # ok
handle_numbers(range(69)) # ok
Let's modify our adopt_kittens method:
from collections.abc import Iterable
class Cat:
...
def adopt_kittens(self, kittens: "Iterable[Cat]") -> None:
self._kittens.extend(kittens)
Siblings of Iterable
Here are some other items from collections.abc you'll find useful:
| Item | Description |
|---|---|
Iterator[T] |
Something you can call next() on. See iterator |
Mapping[K, V] |
An object supporting mapping[key] syntax and other read-only dict-like operations (like mapping.items()). See mapping |
Sequence[T] |
An object supporting iteration, indexing, len() and in. list, tuple and range are all sequences |
MutableSequence[T] |
Like Sequence, but also supports mutation with obj[index] = value, clear(), append() etc. |
Collection[T] |
An object supporting iteration, len() and in. set is a collection but not a sequence |
Callable |
Something you can call(). We will discuss it in a later chapter |
To solve the current problems with kittens(), let's use Collection instead of list:
from collections.abc import Collection
class Cat:
...
def kittens(self) -> "Collection[Cat]":
return self._kittens
Data classes
Classes are commonly used to bundle related pieces of information:
class Dog:
def __init__(self, age: int, name: str, height: float, length: float) -> None:
self.age = age
self.name = name
self.height = height
self.length = length
def cross_section(self) -> float:
return self.height * self.length
def __eq__(self, other: object) -> bool:
if not isinstance(other, Dog):
return NotImplemented
return (
(self.age, self.name, self.height, self.length)
== (other.age, other.name, other.length, other.height)
)
def __repr__(self) -> str:
return (
"Dog("
f"age={self.age!r}"
f", name={self.name!r}"
f", height={self.height!r}"
f", self.lenght={self.length!r}"
")"
)
This is rather verbose, repetitive, and error-prone. It could be much easier to read and write.
The standard library has a dataclasses module:
from dataclasses import dataclass
@dataclass
class Dog:
age: int
name: str
height: float
length: float
def cross_section(self) -> float:
return self.height * self.length
This does exactly the same as the previous snippet (Except for the typos. Did you spot them?).
>>> dog = Dog(age=1, name="fido", height=100.0, length=60.0)
>>> dog
Dog(age=1, name='fido', height=100.0, length=60.0)
>>> dog.name
'fido'
>>> dog.name = 'treat enjoyer'
>>> dog
Dog(age=1, name='treat enjoyer', height=100.0, length=60.0)
>>> dog.cross_section()
6000.0
@dataclass has several flags. We recommend using these in most cases:
kw_onlyprevents you from calling the class with positional arguments, like It's not immediately clear what1,100.0and60.0are, so keyword-only arguments make sense here.frozenprevents you from mutating the object after creation.