Functions Returning Dataclasses

from dataclasses import dataclass
from bakery import assert_equal

@dataclass
class Card:
    suit: str
    rank: int

def parse_card(card: str) -> Card:
    return Card(card[0], int(card[1:]))

assert_equal(parse_card("H2"), Card("H", 2))
assert_equal(parse_card("D5"), Card("D", 5))

Functions not only can consume dataclasses, but they can also create them. In this example, we have a parse_card function that consumes a string value and produces a Card instance. To do so, the function must call the Card constructor and pass in the appropriate fields. In this case, that means the first character of the string as the first argument, and then the rest of the string converted to an integer for the second field. In the assert_equal, we call the function and then assert that the result is the same as calling the constructor with the literal values. In this way, we have a function that can return a dataclass — specifically, a new instance of the dataclass. The utility of a Dataclass is that we can bundle related fields together to represent a more complex object and pass them conveniently around the program without having to have a huge number of variables.

Mutable Updates

from bakery import assert_equal
from dataclasses import dataclass
@dataclass
class Rectangle:
    length: int
    width: int

def stretch(rect: Rectangle) -> Rectangle:
    rect.length = rect.length * 2
    return rect

box = Rectangle(5, 3)
assert_equal(stretch(box), Rectangle(10, 3))

One of the most useful but confusing aspects of a dataclass is that fields’ values can also be changed later. In the example shown here, the stretch function consumes and produces the same Rectangle. However, before returning the Rectangle, its length field is doubled from its previous value. The syntax for updating the fields of a dataclass is almost the same as updating the value of a variable.

Mutability and the Stack/Heap

A stack/heap diagram where the function is modifying the instance passed in.

What do we mean by “the same Rectangle”? This Stack/Heap diagram gives some idea. The stretch function has its own rect parameter, but the instance is the same as the instance pointed to by the box variable in the global scope. The return value comes from the rect instance, so the wide_box variable also ends up holding the same instance. So, the function manipulates the dataclass coming in instead of creating a new one. As we will discuss later in the section about mutability, there are many dangers to mutably updating fields like this. Therefore, we will look at an alternative.

Immutable Copies

from bakery import assert_equal
from dataclasses import dataclass

@dataclass
class Rectangle:
    length: int
    width: int

def stretch(rect: Rectangle) -> Rectangle:
    new_rect = Rectangle(rect.length*2, rect.width)
    return new_rect

box = Rectangle(5, 3)
print(stretch(box))
print(box)

A common way to avoid mutation is to create a brand-new instance with new values instead of updating the old instance. This may sound wasteful, but fortunately Python is quite fast at creating new instances. And in turn, we can use the new instance without worrying about the original being affected. Consider this new version of stretch that creates a new Rectangle instance using the constructor instead of returning the original rect.

Immutability and the Stack/Heap

A stack/heap diagram of a function NOT modifying the instance that was passed in.

Compare the diagram from before with the following new diagram. In the mutable version, there was only one Rectangle instance, which ended up being pointed to by all the variables. But in this version, there’s a second Rectangle instance created during the stretch function and stored in its new_rect local variable. That instance is returned, so the wide_box global variable points to the longer Rectangle instance. The original box variable was not overwritten, and the original Rectangle was not modified.

Under what circumstances would you want to modify the original as opposed to making a copy? The tradeoff might not be clear just yet, but for now, remember that there are two ways to change a dataclass: by making a new one based on an old one, or by editing the fields of an existing one.

Converting Dataclasses

from dataclasses import dataclass
from bakery import assert_equal

@dataclass
class Square:
    color: str
    width: int

@dataclass
class Circle:
    radius: int
    color: str

def square_to_circle(a_square: Square) -> Circle:
    return Circle(a_square.width, a_square.color)

assert_equal(square_to_circle(Square("red", 5)), Circle(5, "red"))

Nothing says that the dataclass coming in must be the same type of dataclass being returned. You might think of a function that takes in one type of dataclass and returns another as a “conversion” function, just like the int and str functions that we have seen before. However, these kinds of functions are often simpler than they might seem from the outside.

The square_to_circle function takes in a Square and returns a Circle. To do so, the function must access the attributes of the a_square parameter and pass them as arguments to the Circle constructor. The body does not have any magic code that “converts” the entire a_square instance all at once; instead, it unpacks the fields into the constructor function. Those fields could be manipulated in more complicated expressions, but fundamentally the process is just like creating any other instance of a dataclass.

Summary