Use more descriptive names for your attributes.
The structure of your data is clear: all members of the Unit
enumeration have values that are 2-item tuples, the first item of which is the unit's symbol, and the second item of which is the unit's domain. The names for Unit
's attributes should reflect this, rather than txt
and col
, both of which are opaque to anyone else reading your code.
(col
looks suspiciously like you've named this attribute according to how you intend to use your Unit
enumeration in your pandas
dataframe. Don't do that. Your Unit
class should, ideally, make sense as an abstraction in its own terms, without reference to any external functions or classes. It should know nothing about how you intend to use it in your wider code base.)
Moreover, it's unclear why you have default values in your __init__
signature, given that each Unit
member must (and does) have a two-item tuple as its value.
You should refactor your __init__
and __str__
methods to the following:
from enum import Enum
class Unit(Enum):
# <-- snip -->
def __init__(self, symbol: str, domain: str) -> None:
self._symbol = symbol
self._domain = domain
# <-- snip -->
def __str__(self) -> str:
return self._symbol
Give your enum members more descriptive names!
Currently, you have duplication of data in your enumeration: each member's name is the same as the first item of its value, the only difference being that the members all have all-uppercase names, while not all members have all-uppercase symbol
s.
You can make your code less repetitive and more readable by giving your enum members more descriptive names:
from enum import Enum
class Unit(Enum):
HOUR = "h", "duration"
MEGAWATT = "MW", "power"
MEGAWATT_HOUR = "MWh", "energy"
EURO = "Eur", "revenue"
EURO_PER_MEGAWATT_HOUR = "Eur/MWh", "price"
DEGREES_CENTIGRADE = "degC", "temperature"
# <-- snip -->
... You probably shouldn't have an __init__
method at all.
The pattern you're implementing in your code is essentially a NamedTuple
/Enum
hybrid. You have structured data, much like you would with a series of NamedTuple
s all of the same type (each member's value is a 2-item tuple; the first item of each member's value is always the unit's symbol, while the second item is always the unit's domain). Meanwhile, you have a predefined set of Unit
s that is known at compilation time and cannot be extended at runtime, much like the traditional conception of an Enum
.
By using an __init__
method to name the fields in your Enum
members' values, however, you in fact add mutable attributes to your enumeration's members. Say we have the following definitions:
from typing import NamedTuple
from enum import Enum
class UnitNT(NamedTuple):
symbol: str
domain: str
class UnitEnum(Enum):
def __init__(self, symbol: str, domain: str) -> None:
self._symbol = symbol
self._domain = domain
HOUR = 'h', 'duration'
Let's observe how these behave in the interactive REPL. Here's the NamedTuple
version:
>>> hour_nt = UnitNT(symbol='h', domain='duration')
>>> hour_nt
UnitNT(symbol='h', domain='duration')
>>> hour_nt.symbol
'h'
>>> hour_nt.symbol = 'foo'
Traceback (most recent call last):
File "<string>", line 1, in <module>
AttributeError: can't set attribute
>>> hour_nt.symbol
'h'
>>> hour_nt
UnitNT(symbol='h', domain='duration')
And here's the Enum
version:
>>> hour_enum = UnitEnum.HOUR
>>> hour_enum
<UnitEnum.HOUR: ('h', 'duration')>
>>> hour_enum._symbol
'h'
>>> hour_enum.value
('h', 'duration')
>>> hour_enum._symbol = 'foo'
>>> hour_enum._symbol
'foo'
>>> hour_enum
<UnitEnum.HOUR: ('h', 'duration')>
>>> hour_enum.value
('h', 'duration')
As we can see, the NamedTuple
version here is truly immutable, and raises an exception if we try to alter the .symbol
attribute of the data, leaving the original data unchanged. However, the behaviour of the Enum
is... surprising. If we try to change the ._symbol
attribute... no error is raised! The ._symbol
attribute is a mutable attribute, just like you'd get in a normal python class. But when we examine the enum member again, we find that, although the ._symbol
attribute has been altered, the member's value is unchanged. All that we've done is created a situation in which the ._symbol
attribute no longer corresponds to the first item of the member's value.
I don't consider this a bug in the design of python enum
s, as there are some situations in which you might want to attach a mutable attribute to an otherwise-immutable constant.
(For example: consider a pack of cards in a pygame
game. There will only ever be 52 cards, so an Enum
might make sense as a good way of representing the pack. Nonetheless, it may be desirable for each member of the pack to have mutable attributes describing the location of the card on the screen, etc., as well as immutable attributes that will never change — the card's suit, the card's rank, etc.)
However, it's undesirable behaviour in this situation. You should get rid of your __init__
method and replace it with read-only properties. While we're at it, we can also improve our __repr__
method to make it more NamedTuple
-ish.
class Unit(Enum):
HOUR = "h", "duration"
MEGAWATT = "MW", "power"
MEGAWATT_HOUR = "MWh", "energy"
EURO = "Eur", "revenue"
EURO_PER_MEGAWATT_HOUR = "Eur/MWh", "price"
DEGREES_CENTIGRADE = "degC", "temperature"
@property
def symbol(self) -> str:
"""Get the symbol which is most commonly used for this unit."""
return self.value[0]
@property
def domain(self) -> str:
"""Get the domain for which this unit is relevant."""
return self.value[1]
# __mul__ and __truediv__ skipped (for now)
def __str__(self) -> str:
return self.symbol
def __repr__(self) -> str:
return f'<Unit.{self.name}(symbol={self.symbol!r}, domain={self.domain!r})>'
Now we have the best parts of NamedTuple
and Enum
combined in our class.
Your __mul__
and __truediv__
methods are not extensible or self-documenting.
In these two methods, you define certain groups of units for which multiplication/division is defined. However, as you note, this logic is repeated between the two methods. Moreover, it's not great having this data hardcoded into the middle of this method at all. This information is fundamental to the definition of the enum, so should be in the class namespace rather than buried in a method.
I'd refactor your code like so:
from __future__ import annotations
from enum import Enum
from typing import NamedTuple
from functools import cache
class MultipliableUnitGroup(NamedTuple):
"""
A class defining a relationship between two `Unit` enum members
such that they can be multiplied together to create a third unit.
"""
multipliers: frozenset[Unit]
result_unit: Unit
class Unit(Enum):
HOUR = "h", "duration"
MEGAWATT = "MW", "power"
MEGAWATT_HOUR = "MWh", "energy"
EURO = "Eur", "revenue"
EURO_PER_MEGAWATT_HOUR = "Eur/MWh", "price"
DEGREES_CENTIGRADE = "degC", "temperature"
@property
def symbol(self) -> str:
"""Get the symbol which is most commonly used for this unit."""
return self.value[0]
@property
def domain(self) -> str:
"""Get the domain for which this unit is relevant."""
return self.value[1]
@classmethod
@property
@cache
def multipliable_unit_groups(cls) -> frozenset[MultipliableUnitGroup]:
"""Get the subgroups of members for which multiplication and division are defined."""
return frozenset({
MultipliableUnitGroup(
multipliers=frozenset({cls.MEGAWATT, cls.HOUR}),
result_unit=cls.MEGAWATT_HOUR
),
MultipliableUnitGroup(
multipliers=frozenset({cls.EURO_PER_MEGAWATT_HOUR, cls.MEGAWATT_HOUR}),
result_unit=cls.EURO
)
})
def __mul__(self, other: Unit) -> Unit:
if type(other) is Unit:
as_set = {self, other}
for multipliers, result_unit in self.multipliable_unit_groups:
if multipliers == as_set:
return result_unit
return NotImplemented
__rmul__ = __mul__
def __truediv__(self, other: Unit) -> Unit:
if type(other) is Unit:
for multipliers, result_unit in self.multipliable_unit_groups:
if result_unit is self and other in multipliers:
return next(filter(other.__ne__, multipliers))
return NotImplemented
def __rtruediv__(self, other: Unit) -> Unit:
return (other / self) if type(other) is Unit else NotImplemented
def __str__(self) -> str:
return self.symbol
def __repr__(self) -> str:
return f'<Unit.{self.name}(symbol={self.symbol!r}, domain={self.domain!r})>'
Since python 3.9, we've been able to stack @classmethod
on top of @property
, which is extremely helpful if we want to add read-only class attributes to an Enum
that we don't want to be converted into members. By throwing functools.cache
into the mix as well, we ensure that the class attribute is only computed once.
Note also that I changed your isinstance
check to an if type(other) is Unit
check — since enums that have members are not subclassable, this makes more sense.
Lastly, it's good practice to return NotImplemented
for undefined operations in methods where you're overloading an operator, rather than raising NotImplementedError
. This is because the object on the right-hand side of the operator might know how to multiply the two objects together, even if the object on the left side doesn't. In the code x * y
, python will first try multiplying the two items together by using x.__mul__
, but if that returns NotImplemented
, it will try again using y.__rmul__
, and only if that also returns NotImplemented
will it then raise a TypeError
telling you that that operation is undefined between x
and y
due to incompatible types. If x.__mul__
raises NotImplementedError
instead of returning NotImplemented
, however, python has no opportunity to try using y.__rmul__
; the raised exception means the programme has already ended. For the same reason, you should always define __rmul__
and __rtruediv__
whenever you define __mul__
and __truediv__
in a class, as I've done above.