feat: algebraic python enums (draft 4)

2026-06-02 05:41:17 -03:00 · 2025-11-03 08:21:38 +05:30
parent c7f25090c4
commit 3eb90bb21c
1 changed files with 189 additions and 19 deletions
--- a/content/post/algebraic-python-enums.md
+++ b/content/post/algebraic-python-enums.md
@@ -9,15 +9,25 @@ tags:
 - Algebraic Data Types
 ---

-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.
+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.
+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](https://peps.python.org/pep-0636/).
+Although Python has the answer to creating structs as
+[dataclasses](https://peps.python.org/pep-0557/), including support for
+[structural match expressions](https://peps.python.org/pep-0636/) in recent
+versions, most tutorials will suggest `Union` types as the equivalent to Rust's
+enums.

-To that end, creating the equivalent to Rust's enum types involves Python union types.
+> 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.

-## First draft
+## Naive draft
+
+{{< collapsable-explanation >}}

 ```python
 # glass_enum.py
@@ -46,17 +56,110 @@ def report_drink(glass: Glass) -> str:
      return f"Ah a {drink}, what a fine taste!"
 ```

+For example
+
+```python
+dr_pepper = Full('Dr. Pepper')
+print(report_drink(dr_pepper))
+```
+
+```python
+# 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))
+```
+
+{{< /collapsable-explanation >}}
+
+Will output
+
+```
+Ah a Dr. Pepper, 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.
+{{< collapsable-explanation >}}
+
+What if we had another `Union` with same variant names in the same file?
+
+```python
+@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`.
+
+```python
+lemonade = Full("lemonade")
+```
+
+```python
+@dataclass
+class Full:
+  gold: int
+  gems: int
+
+@dataclass
+class Empty:
+  pass
+
+Inventory = Full | Empty
+player_inventory = Full(500, 50)
+lemonade = Full("lemonade")
+```
+
+{{< /collapsable-explanation >}}
+
+
+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.

 ```python
 dr_pepper = Glass.Full("Dr. Pepper")
 ```

-This can be partially solved by putting the entire enumerable type inside a module.
+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`.

@@ -76,14 +179,16 @@ 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.
+Since module namespacing 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.
-
+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.
+
 ```python
 def refill(glass: Glass) -> Glass:
  if glass.is_empty(): # can't implement on type `Glass` directly
@@ -91,10 +196,8 @@ def refill(glass: Glass) -> Glass:
  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.
+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

@@ -125,6 +228,8 @@ If only we could register the variants as the `Glass` type itself and inherit al

 We can define a decorator that takes all of the nested dataclasses and makes them inherit the outer class.

+{{< collapsable-explanation >}}
+
 ```python
 import inspect

@@ -136,6 +241,9 @@ def AlgebraicEnum(cls):
    return cls
 ```

+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`.

@@ -152,14 +260,14 @@ class Glass:
  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
@@ -173,11 +281,73 @@ compare it with `Glass.Empty` and `Glass.Full`.

 These methods get automatically called via the method resolution order chain due to the inheritance.

+The following code runs just fine.
+
+```python
+diet_coke = Glass.Full('diet coke')
+empty = Glass.Empty()
+
+print(diet_coke.report_drink())
+print(empty.is_empty())
+```
+
+```python
+import inspect
+from dataclasses import dataclass
+
+
+def AlgebraicEnum(cls):
+    for subclass_name, subclass in inspect.getmembers(cls, predicate=inspect.isclass):
+        if subclass_name != "__class__":
+            setattr(cls, subclass_name, type(subclass_name, (cls, subclass), {}))
+
+    return cls
+
+
+@AlgebraicEnum
+class Glass:
+  @dataclass
+  class Empty:
+    pass
+
+  @dataclass
+  class Full:
+    drink: 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: 'Glass') -> bool:
+    match self:
+      case Glass.Empty():
+        return True
+    return False
+
+
+diet_coke = Glass.Full('diet coke')
+empty = Glass.Empty()
+
+print(diet_coke.report_drink())
+print(empty.is_empty())
+```
+
+{{< /collapsable-explanation >}}
+
+```
+Ah a diet coke, what a fine taste!
+True
+```
+
 ## 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.
+Those 6 lines are the bare minimum to get well organized and namespaced
+algebraic enums in Python that are somewhat comparable to the ones in Rust. These
+enums also play nicely with static type checkers, _goto-definitions_ will
+lead you to the correct class definition.

 I have packaged this decorator with a couple more typing restrictions into a library at [github:lavafroth/ape](https://github.com/lavafroth/ape).