#! Ramblings
of an autodidact...
#! Ramblings

More advanced Dataclass sorting

in PythonEstimated read time: 4 min

Share Tweet Share

A useful tip for making your dataclasses more useful

More advanced Dataclass sorting

While checking my email this morning I came across another useful Pybites tip.

dataclass-order

It was a great little tip, but I wondered, "what if I wanted to sort by bite level?"

Pretty up the output

I know it's just an example code, but I like to make my output a bit more readable. A simple loop of the results is all it takes:

for bite in sorted(bites):
    print(bite)

Output:

Bite(number=11, title='Enrich a class', level=<BiteLevel.ADVANCED: 4>)
Bite(number=21, title='Query a nested DS', level=<BiteLevel.BEGINNER: 2>)
Bite(number=25, title='No promo twice', level=<BiteLevel.INTERMEDIATE: 3>)

Ah, much better!

Using a sort_index

As you can see, the bites are being sorted by their number, but I want to see them sorted by difficulty. The first thing I did was add a sort_index to the dataclass and set it to the property that I wanted to sort by, level:

@dataclass(order=True)
class Bite:
    sort_index: int
    number: int
    title: str
    level: int = BiteLevel.BEGINNER

    def __post_init__(self):
        self.sort_index = self.level

Output:

Traceback (most recent call last):
  File "enum_play.py", line 28, in <module>
    Bite(21, "Query a nested DS"),
TypeError: __init__() missing 1 required positional argument: 'title'

Uh oh! Although it's complaining about a missing title, the problem is actually that it's looking for the sort_index.

Dataclasses field(init=False)

I don't want to initialize that variable, so I make the following change to the Bite class:

from dataclasses import dataclass, field
...

class Bite:
    sort_index: int = field(init=False)
    number: int
    ...

We first import field from dataclasses and then set it to init=False in order to prevent that variable from being initialized.

Bite(sort_index=<BiteLevel.BEGINNER: 2>, number=21, title='Query a nested DS', level=<BiteLevel.BEGINNER: 2>)
Bite(sort_index=<BiteLevel.INTERMEDIATE: 3>, number=25, title='No promo twice', level=<BiteLevel.INTERMEDIATE: 3>)
Bite(sort_index=<BiteLevel.ADVANCED: 4>, number=11, title='Enrich a class', level=<BiteLevel.ADVANCED: 4>)

Um, ok better but I don't want to see sort_index being displayed in the output.

Dataclasses field(repr=False)

In order to prevent sort_index from being displayed, we will have to set the repr to False.

class Bite:
    sort_index: int = field(init=False, repr=False)
    ...

Output:

Bite(number=21, title='Query a nested DS', level=<BiteLevel.BEGINNER: 2>)
Bite(number=25, title='No promo twice', level=<BiteLevel.INTERMEDIATE: 3>)
Bite(number=11, title='Enrich a class', level=<BiteLevel.ADVANCED: 4>)

Much better!

Dataclasses frozen

Looking much better! But what if I wanted to make the dataclass so that it could not be changed in the code? That's where frozen comes in:

@dataclass(order=True, frozen=True)
class Bite:
    sort_index: int = field(init=False, repr=False)
    ...

Output:

Traceback (most recent call last):
  File "enum_play.py", line 25, in <module>
    Bite(11, "Enrich a class", BiteLevel.ADVANCED),
  File "<string>", line 6, in __init__
  File "enum_play.py", line 21, in __post_init__
    self.sort_index = self.level
  File "<string>", line 4, in __setattr__
dataclasses.FrozenInstanceError: cannot assign to field 'sort_index'

Ugh, well at least we know it works!

Setting class variables with setattr

It's complaining because we are setting the index_value in __post_init__. In order to get around it, we'll have to set that value a different way.

...

    def __post_init__(self):
        object.__setattr__(self, "sort_index", self.level)

Output:

Bite(number=21, title='Query a nested DS', level=<BiteLevel.BEGINNER: 2>)
Bite(number=25, title='No promo twice', level=<BiteLevel.INTERMEDIATE: 3>)
Bite(number=11, title='Enrich a class', level=<BiteLevel.ADVANCED: 4>)

There we go!

Conclusion

As you can see, working with dataclasses isn't that hard at all and will save you from having to type too much boilerplate code.

Final code:

from dataclasses import dataclass, field
from enum import IntEnum


class BiteLevel(IntEnum):
    BEGINNER = 2
    INTERMEDIATE = 3
    ADVANCED = 4


@dataclass(order=True, frozen=True)
class Bite:
    sort_index: int = field(init=False, repr=False)
    number: int
    title: str
    level: int = BiteLevel.BEGINNER

    def __post_init__(self):
        object.__setattr__(self, "sort_index", self.level)


bites = [
    Bite(11, "Enrich a class", BiteLevel.ADVANCED),
    Bite(21, "Query a nested DS"),
    Bite(25, "No promo twice", BiteLevel.INTERMEDIATE),
]

for bite in sorted(bites):
    print(bite)

If you would like to learn more about Dataclasses, check out the documentation.

Receive Updates

ATOM

Contacts