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 symbols.
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 NamedTuples 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 Units 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 enums, 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.
