Write Python like you write Rust!

I started programming in Rust a few years ago, and it gradually changed the way I design programs in other programming languages, especially Python. Before I started using Rust, I usually wrote Python code in a very dynamic and loosely typed way, with no type hints, passing and returning dictionaries here and there, with the occasional fallback to “string-typed” interfaces. However, after experiencing the rigor of Rust’s type system, and noticing all the problems it prevents “by construction”, I suddenly become very anxious whenever I go back to Python and don’t get the same guarantees.

To be clear, the “guarantee” here does not mean memory safety (Python itself is reasonably memory safe), but “robustness” – the concept of designing APIs that are difficult or impossible to abuse, so as to prevent unintended Defined behavior and various errors. In Rust, incorrectly used interfaces often result in compilation errors. In Python you can still execute such incorrect programs, but if you use a type checker like pyright or an IDE with a type analyzer like PyCharm you can still get a similar level of insight about possible problems Quick feedback.

Eventually, I started to adopt some of Rust’s concepts in my Python programs. It basically boils down to two things – use type hints as much as possible, and stick to the principle of making illegal state unrepresentable. I try to do this both for programs that will be maintained for a while and for oneshot utility scripts. Mainly because in my experience the latter often turns into the former ๐Ÿ™‚ In my experience this approach results in programs that are easier to understand and change.

In this article, I will show several examples of such patterns applied to Python programs. It’s not rocket science, but I still think it might be useful to document them.

Note: This article contains a lot of perspective on writing Python code. I don’t want to put “IMHO” in every sentence, so take everything in this post just as my opinion on the matter, not trying to promote some general truth ๐Ÿ™‚ Also, I’m not saying The idea presented is that all of these were invented in Rust, and of course, they are used in other languages as well.

Type hint

The first and foremost is to use type hints wherever possible, especially in function signatures and class attributes. When I read a function signature like this:

def find_item(records, check):

I don’t know what happened to the signature itself. Is it a list of records, a dictionary or a database connection? Is check a boolean or a function? What does this function return? What happens if it fails, does it raise an exception or return None? To find answers to these questions, I’d either have to read the function body (and often recursively read the function bodies of other functions it calls – which is annoying), or read its documentation (if there is one). While documentation may contain useful information about what a function does, it is not necessary to use it to also document answers to previous questions. Many questions can be answered by a built-in mechanism – type hints.

def find_item(
  records: List[Item],
  check: Callable[[Item], bool]
) -> Optional[Item]:

Did I spend more time writing my signature? Yes. Is that the problem? No, unless my encoding is bottlenecked by the number of characters written per minute, which doesn’t really happen. Writing out the type explicitly forced me to think about what the actual interface the function provides, and how to make it as strict as possible so that it’s hard for its callers to use it in the wrong way. With the signature above, I can get a good idea of how to use the function, what to pass to it as parameters, and what I expect to return from it. Also, unlike doc comments, which are easily outdated when the code changes, the type checker yells at me when I change types and doesn’t update the callers of the function. If I’m interested in what an Item is, I can just go to definition and see what that type looks like right away.

I’m not an absolutist in this regard, and if it takes five nested type hints to describe a single parameter, I usually give up and give it a simpler but imprecise type. In my experience, this doesn’t happen very often. If it does happen, it might actually indicate something wrong with your code – if your function argument can be a number, a tuple of strings, or a dictionary mapping strings to integers, this might be an indication that you might want to re- Construct and simplify it.

Dataclass instead of tuple or dictionary

Using type hints is one thing, but that only describes what the function’s interface is. The second step is actually to make these interfaces as precise and “locked down” as possible. A typical example is returning multiple values (or a complex value) from a function. The lazy and fast way is to return a tuple:

def find_person(...) -> Tuple[str, str, int]:

Great, we know we’re going to return three values. what are these? Is the first string the person’s name? Second string last name? what is your phone number? Is it age? position in some list? social Security number? This kind of input is opaque, and you don’t know what’s going on here unless you look at the function body.

The next step to “improve” this might be to return a dictionary:

def find_person(...) -> Dict[str, Any]:
    ...
    return {
        "name": ...,
        "city": ...,
        "age": ...
    }

Now we actually know what each returned property is, but we have to inspect the function body again to find out. In a sense, typing gets worse because now we don’t even know the number and type of the individual properties. Also, when this function changes and a key in the returned dictionary is renamed or removed, there’s no easy way to find out with the type checker, so its callers often have to use very manual and annoying run-crash- Modify the code to change the loop.

The correct solution is to return a strongly typed object whose named parameters have additional types. In Python, this means we have to create a class. I suspect tuples and dictionaries are often used in these cases, since it’s much easier than defining a class (and giving it a name), creating a constructor with parameters, storing parameters to fields, etc. Since Python 3.7 (and faster using the package polyfill), there is a faster solution – dataclasses.

@dataclasses.dataclass
class City:
    name: str
    zip_code: int


@dataclasses.dataclass
class Person:
    name: str
    city: City
    age: int


def find_person(...) -> Person:

You still have to think about a name for the class you create, but other than that, it’s as concise as possible, and you get type annotations for all properties.

With this data class, I have a clear description of what the function returns. When I call this function and process the return value, the IDE autocomplete will show me the name and type of its property. This might sound trivial, but to me it’s a huge productivity advantage. Also, when code is refactored and properties change, my IDE and type checker will yell at me and show me all the places where I have to change without me having to execute the program at all. For some simple refactorings (such as property renaming), the IDE even makes these changes for me. Also, with explicitly named types, I can build a vocabulary of terms (Person,City) that can then be shared with other functions and classes.

Algebraic data types

The thing I’m probably most lacking in Rust among most mainstream languages is Algebraic Data Types (ADT)2. It’s a very powerful tool for explicitly describing the shape of the data my code is working with. For example, when I’m working with packets in Rust, I can explicitly enumerate all the various packets that can be received, and assign each of them different data (fields):

enum Packet {
  Header {
    protocol: Protocol,
    size: usize
  },
  Payload {
    data: Vec<u8>
  },
  Trailer {
    data: Vec<u8>,
    checksum: usize
  }
}

With pattern matching, I can react to individual variants, and the compiler checks that I haven’t missed any cases:

fn handle_packet(packet: Packet) {
  match packet {
    Packet::Header { protocol, size } => ...,
    Packet::Payload { data } |
    Packet::Trailer { data, ...} => println!("{data:?}")
  }
}

This is invaluable for ensuring that invalid state is not representable and thus avoiding many runtime errors. ADTs are especially useful in statically typed languages, where if you want to use a set of types in a uniform way, you need a shared “name” to refer to them. Without an ADT, this is usually done using OOP interfaces and/or inheritance. Interfaces and virtual methods have their place when the set of types used is open, but when the set of types is closed and you want to make sure you handle all possible variants, ADT and pattern matching are more appropriate.

In a dynamically typed language like Python, there’s actually no need to share a name for a set of types, mostly because you don’t even have to name the types you use in your program in the first place. However, it is still useful to use something like an ADT by creating a union type:

@dataclass
class Header:
  protocol: protocol
  size: int

@dataclass
class Payload:
  data: str

@dataclass
class Trailer:
  data: str
  checksum: int

Packet = typing. Union[Header, Payload, Trailer]
# or `Packet = Header | Payload | Trailer` since Python 3.10

Packet defines a new type here, which can be a header, payload, or trailer packet. I can now use this type (name) in the rest of the program when I want to make sure only these three classes are valid. Note that classes have no explicit “label” attached to them, so when we want to differentiate them we have to use eginstanceof or pattern matching:

def handle_is_instance(packet: Packet):
    if isinstance(packet, Header):
        print("header {packet. protocol} {packet. size}")
    elif isinstance(packet, Payload):
        print("payload {packet.data}")
    elif isinstance(packet, Trailer):
        print("trailer {packet. checksum} {packet. data}")
    else:
        assert False

def handle_pattern_matching(packet: Packet):
    match packet:
        case Header(protocol, size): print(f"header {protocol} {size}")
        case Payload(data): print("payload {data}")
        case Trailer(data, checksum): print(f"trailer {checksum} {data}")
        case _: assert False

Sadly, here we have to (or rather, should) include the annoying assert False branch so that the function crashes when it receives unexpected data. In Rust, this would be a compile-time error.

Note: Several people on Reddit have alerted me that assert False actually completely optimizes away python -O … in an optimized build ( ). Therefore, it is safer to throw the exception directly. There’s also typing.assert_never from Python 3.11, which explicitly tells the type checker that falling into this branch should be a “compile-time” error.

A nice property of a union type is that it is defined outside of the class that is part of the union. So the class doesn’t know it’s contained in the union, which reduces coupling in the code. You can even create multiple different unions with the same type:

Packet = Header | Payload | Trailer
PacketWithData = Payload | Trailer

Union types are also very useful for automatic (de)serialization. Recently I discovered a great serialization library called pyserde which is based on the old Rustserde serialization framework. Among many other cool features, it is able to leverage type annotations to serialize and deserialize union types without any extra code:

import serde

...
Packet = Header | Payload | Trailer

@dataclass
class Data:
    packet: Packet

serialized = serde.to_dict(Data(packet=Trailer(data="foo", checksum=42)))
# {'packet': {'Trailer': {'data': 'foo', 'checksum': 42}}}

deserialized = serde.from_dict(Data, serialized)
# Data(packet=Trailer(data='foo', checksum=42))

You can even choose how union tags are serialized, with serde. I was looking for a similar feature, as it is very useful for (de)serializing union types. dataclasses_json however, is very annoying to implement in most other serialization libraries I’ve tried (eg or dacite.

For example, when working with machine learning models, I use federation to store various types of neural networks (such as classification or segmentation CNN models) in a single configuration file format. I’ve also found it useful to version data in different formats (config files in my case), like so:

Config = ConfigV1 | ConfigV2 | ConfigV3

By deserializing Config, I was able to read all previous versions of the config format, thus maintaining backwards compatibility.

use newtype

In Rust, it is common to define data types that do not add any new behavior, but are only used to specify the domain and intended use of some other very general data types, such as integers. This mode is called “newtype”3, and it can also be used in Python. Here’s a motivating example:

class Database:
  def get_car_id(self, brand: str) -> int:
  def get_driver_id(self, name: str) -> int:
  def get_ride_info(self, car_id: int, driver_id: int) -> RideInfo:

db = Database()
car_id = db.get_car_id("Mazda")
driver_id = db. get_driver_id("Stig")
info = db.get_ride_info(driver_id, car_id)

found mistake?

The parameter get_ride_info is swapped. There are no type errors because both the car ID and the driver ID are simple integers and thus the correct type, even though the function call is semantically wrong.

We can fix this by defining separate types for different types of IDs using “NewType”:

from typing import NewType

# Define a new type called "CarId", which is internally an `int`
CarId = NewType("CarId", int)
# Ditto for "DriverId"
DriverId = NewType("DriverId", int)

class Database:
  def get_car_id(self, brand: str) -> CarId:
  def get_driver_id(self, name: str) -> DriverId:
  def get_ride_info(self, car_id: CarId, driver_id: DriverId) -> RideInfo:


db = Database()
car_id = db.get_car_id("Mazda")
driver_id = db. get_driver_id("Stig")
# Type error here -> DriverId used instead of CarId and vice-versa
info = db.get_ride_info(<error>driver_id</error>, <error>car_id</error>)

This is a very simple pattern that can help catch hard-to-find bugs. It’s especially useful, for example, if you’re dealing with many different types of IDs (CarId vs DriverId) or certain metrics that shouldn’t be mixed together (Speed vs Lengthvs, etc.). temperature

Using constructors

One of the things I really like about Rust is that it doesn’t have constructors per se. Instead, people tend to use plain functions to create (ideally properly initialize) struct instances. In Python, there is no constructor overloading, so if you need to construct an object in more than one way, someone will cause an __init__ method to have many parameters that are used in different ways for initialization and can’t really be used together.

Instead, I like to create “constructor” functions with explicit names, which make it obvious how and from what data an object is constructed:

class Rectangle:
    @staticmethod
    def from_x1x2y1y2(x1: float, ...) -> "Rectangle":
    
    @staticmethod
    def from_tl_and_size(top: float, left: float, width: float, height: float) -> "Rectangle":

This makes constructing objects much cleaner and does not allow users of the class to pass invalid data when constructing objects (for example by combining y1 and width).

Use type-encoded invariants

Using the type system itself to encode invariants that can only be tracked at runtime is a very general and powerful concept. In Python (and other mainstream languages), I often see classes as big, fluffy balls of mutable state. One source of this confusion is code that tries to keep track of object invariants at runtime. It has to account for many situations that could happen in theory, since the type system doesn’t make them impossible (“if the client has been asked to disconnect, and now someone tries to send it a message, but the socket is still connected”, etc. ).

Client

Here’s a typical example:

class Client:
  """
  Rules:
  - Do not call `send_message` before calling `connect` and then `authenticate`.
  - Do not call `connect` or `authenticate` multiple times.
  - Do not call `close` without calling `connect`.
  - Do not call any method after calling `close`.
  """
  def __init__(self, address: str):

  def connect(self):
  def authenticate(self, password: str):
  def send_message(self, msg: str):
  def close(self):

… easy? You just need to read the documentation carefully and make sure you never violate the above rules (to avoid invoking undefined behavior or crashing). Another approach is to fill the class with various assertions that check all mentioned rules at runtime, which leads to code clutter, missing edge cases, and slower feedback in case of errors (compile time vs. runtime). The heart of the matter is that clients can exist in various (mutually exclusive) states, but instead of modeling these states individually, they are all combined into one type.

Let’s see if we can improve this by splitting the various states into separate types4.

First, does it make sense to have a Client that is not connected to anything? It doesn’t seem to be the case. Such an unconnected client can’t do anything until you call connect anyway. So why allow this state to exist? We can create a constructor called connect that will return a connected client:

def connect(address: str) -> Optional[ConnectedClient]:
  pass

class ConnectedClient:
  def authenticate(...):
  def send_message(...):
  def close(...):

If the function succeeds, it will return a client that supports the “connected” invariant, and you can’t call connect again to screw things up. If the connection fails, the function can raise an exception or return None or some explicit error.

A similar approach can be used for the state authenticated. We can introduce another type that holds the invariant that the client is connected and authenticated:

class ConnectedClient:
  def authenticate(...) -> Optional["AuthenticatedClient"]:

class AuthenticatedClient:
  def send_message(...):
  def close(...):

Only after we actually have an instance of an AuthenticatedClient can we actually start sending messages.

The last problem is the method close. In Rust (due to destructive move semantics), we are able to express the fact that when the close method is called, you can no longer use the client. This is not possible in Python, so we have to use some workarounds. One solution might be to fall back to runtime tracking, introduce a boolean property in the client, and assert close that send_message has not been closed. Another approach might be to remove the close method entirely and just use the client as a context manager:

with connect(...) as client:
    client. send_message("foo")
# Here the client is closed

There is no close method available, and you cannot accidentally close the client twice.

Strongly typed bounding boxes

Object detection is a computer vision task I sometimes work on where the program has to detect a set of bounding boxes in an image. Bounding boxes are basically beautified rectangles with some additional data, and they appear everywhere when you implement object detection. An annoying thing about them is that sometimes they are normalized (the coordinates and size of the rectangle are in the interval [0.0, 1.0]), but sometimes they are denormalized (the coordinates and size are limited by the dimensions of the image they are attached to). When you send a bounding box through many functions that deal with data preprocessing or postprocessing, it’s easy to screw it up, for example by normalizing the bounding box twice, which can lead to very annoying bugs to debug.

This has happened to me a few times, so at one point I decided to fix it completely by separating these two types of bboxes into two different types:

@dataclass
class NormalizedBBox:
  left: float
  top: float
  width: float
  height: float


@dataclass
class DenormalizedBBox:
  left: float
  top: float
  width: float
  height: float

With this separation, normalized and denormalized bounding boxes are no longer easily mixed together, which mainly solves the problem. However, we can make some improvements to make the code more ergonomic:

  • Reduce duplication through composition or inheritance:

@dataclass
class BBoxBase:
  left: float
  top: float
  width: float
  height: float

# Composition
class NormalizedBBox:
  bbox: BBoxBase

class DenormalizedBBox:
  bbox: BBoxBase

Bbox = Union[NormalizedBBox, DenormalizedBBox]

#Inheritance
class NormalizedBBox(BBoxBase):
class DenormalizedBBox(BBoxBase):
  • Add runtime checks to ensure normalized bounding boxes are actually normalized:

class NormalizedBBox(BboxBase):
  def __post_init__(self):
    assert 0.0 <= self. left <= 1.0
    ...
  • Add a method to convert between the two representations. In some places we might want to know the explicit representation, but in other places we want to use the generic interface (“any kind of BBox”). In that case, we should be able to convert “any BBox” into one of two representations:

class BBoxBase:
  def as_normalized(self, size: Size) -> "NormalizeBBox":
  def as_denormalized(self, size: Size) -> "DenormalizedBBox":

class NormalizedBBox(BBoxBase):
  def as_normalized(self, size: Size) -> "NormalizedBBox":
    return self
  def as_denormalized(self, size: Size) -> "DenormalizedBBox":
    return self.denormalize(size)

class DenormalizedBBox(BBoxBase):
  def as_normalized(self, size: Size) -> "NormalizedBBox":
    return self. normalize(size)
  def as_denormalized(self, size: Size) -> "DenormalizedBBox":
    return self

With this interface, I can have the best of both worlds – separate types for correctness, and have a unified interface for ergonomics.

Note: If you want to add some shared methods to parent/base classes that return instances of the corresponding class, you can use typing.Self since Python 3.11:

class BBoxBase:
  def move(self, x: float, y: float) -> typing. Self: ...

class NormalizedBBox(BBoxBase):
  ...

bbox = NormalizedBBox(...)
# The type of `bbox2` is `NormalizedBBox`, not just `BBoxBase`
bbox2 = bbox. move(1, 2)

Safer mutex

Mutexes and locks in Rust are usually provided behind a very nice interface, which has two benefits:

When you lock a mutex, you get a guard object that is automatically unlocked when the mutex is destroyed, using the ancient RAII mechanism:

{
  let guard = mutex. lock(); // locked here
  ...
} // automatically unlocked here

This means you can’t accidentally forget to unlock the mutex. A very similar mechanism is commonly used in C++, although explicit lock/unlock interfaces without guard objects are also available for std::mutex, which means they can still be misused.

Data protected by a mutex is stored directly in the mutex (struct). With this design, it is impossible to access protected data without actually locking the mutex. You must first lock the mutex to get the guard, then use the guard itself to access the data:

let lock = Mutex::new(41); // Create a mutex that stores the data inside
let guard = lock. lock(). unwrap(); // Acquire guard
*guard + = 1; // Modify the data using the guard

This is in stark contrast to the mutex APIs common in mainstream languages (including Python), where the mutex and the data it protects are separate, so you can easily forget to actually lock the mutex before accessing the data:

mutex = Lock()

def thread_fn(data):
    # Acquire mutex. There is no link to the protected variable.
    mutex. acquire()
    data.append(1)
    mutex. release()

data = []
t = Thread(target=thread_fn, args=(data,))
t. start()

# Here we can access the data without locking the mutex.
data.append(2) # Oops

While we don’t get the exact same benefits in Python that we get in Rust, not all is lost. Python locks implement the context manager interface, which means you can use them in blocks with to ensure they are automatically unlocked when the scope ends. With a little effort, we can go further:

import contextlib
from threading import Lock
from typing import ContextManager, Generic, TypeVar

T = TypeVar("T")

# Make the Mutex generic over the value it stores.
# In this way we can get proper typing from the `lock` method.
class Mutex(Generic[T]):
  # Store the protected value inside the mutex
  def __init__(self, value: T):
    # Name it with two underscores to make it a bit harder to accidentally
    # access the value from the outside.
    self.__value = value
    self.__lock = Lock()

  # Provide a context manager `lock` method, which locks the mutex,
  # provides the protected value, and then unlocks the mutex when the
  # context manager ends.
  @contextlib.contextmanager
  def lock(self) -> ContextManager[T]:
    self.__lock. acquire()
    try:
        yield self.__value
    finally:
        self.__lock.release()

# Create a mutex wrapping the data
mutex = Mutex([])

# Lock the mutex for the scope of the `with` block
with mutex. lock() as value:
  # value is typed as `list` here
  value.append(1)

With this design, you can only access protected data after actually locking the mutex. Obviously, this is still Python, so you can still break invariants — for example, by storing another pointer to protected data outside of the mutex. But unless your behavior is hostile, this makes the mutex interface in Python safer to use.

Anyway, I’m sure I’m using more “robust modes” in my Python code, but that’s about all I can think of at the moment. Let me know if you have some examples of similar ideas or any other comments.

  1. To be fair, if you’re using some sort of structured format (like reStructuredText), the parameter type descriptions in the doc comments probably do as well. In that case, the type checker might use it and warn you when the types don’t match. However, if you’re using a type checker anyway, I think it’s best to take advantage of the “native” mechanism for specifying types – type hints.

  2. aka discriminated/tagged unions, sum types, sealed classes, etc.

  3. Yes, there are other use cases for new types besides those described here, stop yelling at me.

  4. This is known as the typestate pattern.

  5. Unless you try hard enough, e.g. to manually call the magic __exit__ method.