More advanced Dataclass sorting
While checking my email this morning I came across another useful Pybites tip.
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.