As much as I like rust for its ergonomic features, University has forced me to use Python for the past couple of months, especially because of the hype for machine learning and data science.
One of the biggest things that I missed from the rust experience was enumerable data types whose variants can wrap around different datatypes.
Fortunately, since Python 3.8, creating structs has been a breeze using the dataclass decorator. There’s even support for structural match expressions, like in rust, in recent versions of Python.
To that end, creating the equivalent to Rust’s enum types involves Python union types.
First draft
# glass_enum.py
+Algebraic Python Enums - lavafroth Algebraic Python Enums
As much as I like Rust for its ergonomic features, university has forced me to
+use Python for the past couple of months, especially because of the hype for
+machine learning and data science.
One of the biggest things that I missed from the rust experience was enumerable
+data types whose variants can wrap around different datatypes.
Although Python has the answer to creating structs as
+dataclasses, including support for
+structural match expressions in recent
+versions, most tutorials will suggest Union types as the equivalent to Rust’s
+enums.
I highly encourage you to try out the code snippets and follow along with this article.
+Use the collapse explanation button to copy multiple code blocks in one go.
Naive draft
# glass_enum.py
from dataclasses import dataclass
@@ -19,8 +34,63 @@ Fortunately, since Python 3.8, creating structs has been a breeze using the data
return "Whoops, looks like you've finished your drink!"
case Full(drink):
return f"Ah a {drink}, what a fine taste!"
-
Pitfalls
No direct variant access
Since there is no namespacing, we also can’t instantiate variants under the Glass namespace. The following code does not work.
dr_pepper = Glass.Full("Dr. Pepper")
-
This can be partially solved by putting the entire enumerable type inside a module.
+
For example
dr_pepper = Full('Dr. Pepper')
+print(report_drink(dr_pepper))
+
# glass_enum.py
+
+from dataclasses import dataclass
+
+@dataclass
+class Empty:
+ pass
+
+@dataclass
+class Full:
+ drink: str
+
+Glass = Empty | Full
+
+def report_drink(glass: Glass) -> str:
+ match glass:
+ case Empty():
+ return "Whoops, looks like you've finished your drink!"
+ case Full(drink):
+ return f"Ah a {drink}, what a fine taste!"
+
+dr_pepper = Full('Dr. Pepper')
+print(report_drink(dr_pepper))
+
Will output
Ah a Dr. Pepper, what a fine taste!
+Pitfalls
No direct variant access
What if we had another Union with same variant names in the same file?
@dataclass
+class Full:
+ gold: int
+ gems: int
+
+@dataclass
+class Empty:
+ pass
+
+Inventory = Full | Empty
+player_inventory = Full(500, 50)
+Now we try to instantiate a Glass Full of lemonade.
lemonade = Full("lemonade")
+@dataclass
+class Full:
+ gold: int
+ gems: int
+
+@dataclass
+class Empty:
+ pass
+
+Inventory = Full | Empty
+player_inventory = Full(500, 50)
+lemonade = Full("lemonade")
+Python will error out since Full now refers to the new variant of the
+union type Inventory.
Traceback (most recent call last):
+ File "<python-input-11>", line 1, in <module>
+ lemonade = Full("lemonade")
+TypeError: Full.__init__() missing 1 required positional argument: 'gems'
+We can’t instantiate variants as members of the Glass namespace. The following code does not work.
dr_pepper = Glass.Full("Dr. Pepper")
+This can be partially solved by keeping just the Glass type inside a module.
Here we have saved the file as glass_enum.py. From a different module we can
access the variants as glass_enum.Empty and glass_enum.Full.
# main.py
import glass_enum
@@ -30,12 +100,14 @@ access the variants as glass_enum.Empty and glass_enum.FullNow any function outside the module has to ingest a rather confusing argument of type glass_enum.Glass.
def refill(glass: glass_enum.Glass) -> glass_enum.Glass:
# ...
return glass
-Since using a module namespace only causes more confusion, we will discard this idea.
No methods on the enum itself
With no namespaceing, methods cannot be defined on the Union of the different variants.
In the case of our concrete example, we can’t add methods to the Glass type.
def refill(glass: Glass) -> Glass:
+Since module namespacing only causes more confusion, we will discard this idea.
No methods on the enum itself
Python also disallows methods from being defined on Union types.
+In the case of our concrete example, we can’t add methods to the Glass type.
The following code uses a hypothetical is_empty() method on the Glass union
+type which is not allowed. Hence the code won’t run.
def refill(glass: Glass) -> Glass:
if glass.is_empty(): # can't implement on type `Glass` directly
return Full('water')
return glass
-Even if we use module level namespacing, it’s simply not possible to define any method on a Union type in Python.
To define a method like is_empty(), it must be implemented on both the classes Full and Empty. This can get
-tedious if there are 3 or more variants.
Python is a sneaky language
Last week I discovered that Python allows creating nested classes to keep things organized.
from dataclasses import dataclass
+To define a method like is_empty(), it must be implemented on both the classes
+Full and Empty. This gets tedious for 3 or more variants.
Python is a sneaky language
Last week I discovered that Python allows creating nested classes to keep things organized.
from dataclasses import dataclass
class Glass:
@dataclass
@@ -46,7 +118,7 @@ tedious if there are 3 or more variants.P
class Full:
drink: str
Python will happily run the above code and we can access the “variants” under the Glass namespace.
lemonade = Glass.Full('lemonade')
-If only we could register the variants as the Glass type itself and inherit all its methods.
Redecorate
We can define a decorator that takes all of the nested dataclasses and makes them inherit the outer class.
import inspect
+If only we could register the variants as the Glass type itself and inherit all its methods.
Redecorate
We can define a decorator that takes all of the nested dataclasses and makes them inherit the outer class.
import inspect
def AlgebraicEnum(cls):
for subclass_name, subclass in inspect.getmembers(cls, predicate=inspect.isclass):
@@ -54,7 +126,8 @@ tedious if there are 3 or more variants.P
setattr(cls, subclass_name, type(subclass_name, (cls, subclass), {}))
return cls
-That’s all there is to the magic! Now we can simply add this decorator above the previous class declaration +
The inheritance means all methods of the outer class are available on the nested
+classes and any object of a nested class isinstance of the outer class.
That’s all there is to the magic! Now we can simply add this decorator above the previous class declaration
and the variants like Glass.Empty and Glass.Full would be of the type Glass.
from dataclasses import dataclass
@AlgebraicEnum
@@ -67,22 +140,71 @@ and the variants like Glass.Empty and Glass.Full would
class Full:
drink: str
- def report_drink(self) -> str:
+ def report_drink(self: 'Glass') -> str:
match self:
case Glass.Empty():
return "Whoops, looks like you've finished your drink!"
case Glass.Full(drink):
return f"Ah a {drink}, what a fine taste!"
- def is_empty(self) -> str:
+ def is_empty(self: 'Glass') -> bool:
match self:
case Glass.Empty():
return True
return False
As a bonus, the variants will also inherit any methods defined on the Glass type.
Note how the report_drink method accepts a self of type Glass and the match arms
-compare it with Glass.Empty and Glass.Full.
These methods get automatically called via the method resolution order chain due to the inheritance.
Closing thoughts
Those 6 lines are the bare minimum of what you can do right now to have well organized and namespaced algebraic enums in Python -which are somewhat comparable to those in Rust. These enums also play nicely with static type checkers -and goto-definitions will also lead you to the correct class defining a variant or the enum itself.
I have packaged this decorator with a couple more typing restrictions into a library at github:lavafroth/ape.
I hope you enjoyed this foray into contorting Python.