15
\$\begingroup\$

I'm using Python's Enum to define the physical units in which a value is expressed.

Eventually, I want to add them as an attribute to my pandas.Series and pandas.DataFrame instances, so that this information does not need to be stored in the variable name or in comments in the code. Also, it allows for automatic calculation of the unit of a calculation, e.g., that the result of 40 megawatts during 5 hours is a value in MWh.

Here's my code; my questions are below.

class Unit(Enum):
    H = ("h", "duration")
    MW = ("MW", "power")
    MWH = ("MWh", "energy")
    EUR = ("Eur", "revenue")
    EURMWH = ("Eur/MWh", "price")
    DEGC = ("degC", "temperature")

    def __init__(self, txt: str = "", col: str = ""):
        self._txt = txt
        self._col = col

    def __mul__(self, other):
        if isinstance(other, Unit):
            relations = ((Unit.MWH, Unit.MW, Unit.H), (Unit.EUR, Unit.MWH, Unit.EURMWH))
            for (prod, mul1, mul2) in relations:
                if (mul1, mul2) == (self, other) or (mul1, mul2) == (other, self):
                    return prod
        raise NotImplementedError("This multiplication is not defined.")

    def __truediv__(self, other):
        if isinstance(other, Unit):
            relations = ((Unit.MWH, Unit.MW, Unit.H), (Unit.EUR, Unit.MWH, Unit.EURMWH))
            for (prod, mul1, mul2) in relations:
                if prod is self and mul1 is other:
                    return mul2
                if prod is self and mul2 is other:
                    return mul1
        raise NotImplementedError("This division is not defined.")

    def __str__(self):
        return self._txt

u1 = Unit.EUR
u2 = Unit.MWH
u3 = u1 / u2 # Unit.EURMWH

Two questions:

  1. The variable relations is repeated. Is there a better way to do this? I'd prefer to keep in inside the class, but I can't put it at the class's root as it'd be interpreted as an enum value. Put it in a @staticmethod and replace references to it by self.relations()?
  2. More generally, is this a valid use case of Enums? It seems sensible to me, but it's going far beyond just being a list of integers as used in other languages.
\$\endgroup\$
0

4 Answers 4

19
\$\begingroup\$

Is this a valid use case of Enums?

In my opinion, using Python's enum module for this kind of task is highly idiomatic — I think you've chosen the best way to express your code's intent here. However, I think there are nonetheless several aspects I might do differently.


  1. 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
    

  2. 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 -->
    

  3. ... 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.


  4. 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.




Addendum

In my final code snippet above, I used two Python 3.9 features: functools.cache and the ability to stack @classmethod on top of @property. However, note that if you're on Python <= 3.8, you can use the following code instead, which is more backwards-compatible. It utilises the fact that methods and properties declared in a metaclass definition have similar behaviour to classmethods and classmethod-properties declared in a class definition:

from __future__ import annotations

from enum import Enum, EnumMeta
from typing import NamedTuple
from functools import lru_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 UnitEnumMeta(EnumMeta):
    """Metaclass for `Unit`"""

    @property
    @lru_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
            )
        })


class Unit(Enum, metaclass=UnitEnumMeta):
    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]        

    def __mul__(self, other: Unit) -> Unit:
        if type(other) is Unit:
            as_set = {self, other}
            for multipliers, result_unit in type(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 type(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})>'
\$\endgroup\$
4
  • 1
    \$\begingroup\$ Great answer @Alex and many good points. Many thanks for taking the time! I was aware of the mutability and had already added property accessors to get the ._txt and ._col values, but forgoing __init__ altogether is a good idea I wasn't aware of. \$\endgroup\$
    – ElRudi
    Sep 13, 2021 at 11:28
  • \$\begingroup\$ @ElRudi No problem - glad I was helpful! Good catch with the suggested edit -- I was editing my answer simultaneously so it got lost, but I've made your suggested change. \$\endgroup\$ Sep 13, 2021 at 11:40
  • 1
    \$\begingroup\$ Thanks also for suggesting return NotImplemented instead of raise NotImplementedError; I wasn't aware of this distinction. About implementing __rmul__: in this particular case, the only implemented case is when 2 units are multiplied, so if __mul__ returns NotImplemented, so does __rmul__. Same for division. Is it still necessary/good practice to define both in that case? (__rmul__ = __mul__ is easy enough, but __rtruediv__ must be implemented seperately.) \$\endgroup\$
    – ElRudi
    Sep 13, 2021 at 12:39
  • \$\begingroup\$ @ELRudi, yes -- it's arguable whether it's strictly necessary to define __rmul__ and __rtruediv__ in this particular case. But I'd say it's a good habit to get into nonetheless. Also, my initial implementation of __rtruediv__ was far more convoluted than it needed to be. You can just delegate to __truediv__ in this situation, making it a one-liner (see my edited version!). \$\endgroup\$ Sep 13, 2021 at 12:42
8
\$\begingroup\$

It's an excellent exercise to try to write your own library.

Since you mentioned pandas and currency units, you might want to try pint.

Pint is a Python package to define, operate and manipulate physical quantities: the product of a numerical value and a unit of measurement. It allows arithmetic operations between them and conversions from and to different units.

You can define your own units and dimensions:

from pint import UnitRegistry

unit = UnitRegistry()

unit.define('dollar = [currency] = USD')
unit.define('euro = 1.18 dollar = EUR')
unit.define('bitcoin = 44733 USD = BTC')

price = 10 * unit.EUR
print(price)
# 10 euro
print(price.to('dollars'))
# 11.799999999999999 dollar
print(price.to('BTC'))
# 0.0002637873605615541 bitcoin
print(10 * unit.EUR + 2 * unit.USD)
# 11.694915254237289 euro

You can now combine monetary units with energy units:

print((123 * unit.EUR / unit.MWh * 5 * unit.kW * 1 * unit.year).to('USD'))
# 6361.486199999999 dollar

Pint seems to support numpy arrays out of the box, so it should work fine with series and dataframes in pandas too:

energy_per_year = [1, 2, 3] * unit.MWh / unit.year
print(energy_per_year)
# [1.0 2.0 3.0] megawatt_hour / year
print(energy_per_year.to(unit.kW))
# [0.11407711613050422 0.22815423226100845 0.3422313483915127] kilowatt

Pint also delivers meaningful error messages for forbidden operations:

print(1 * unit.EUR + 3 * unit.kW)
# pint.errors.DimensionalityError: Cannot convert from 'euro' ([currency]) to 'kilowatt' ([length] ** 2 * [mass] / [time] ** 3)

Since it's open-source, you can take a look at the code and compare it to yours.

\$\endgroup\$
4
  • \$\begingroup\$ I wanted to easily convert time from seconds to months and honestly found pint to be really really easy. +1 One note tho; I found something like unit("20 m / 5 s") to output nonsense. But unit("20 / 5 m / s") seemed to fix the problem. \$\endgroup\$
    – Peilonrayz
    Sep 14, 2021 at 19:20
  • \$\begingroup\$ @Peilonrayz: It's a great library indeed, and the seamless integration with numpy is a nice touch. When writing fractions on a piece of paper, it's really easy to show the difference between \$\frac{ab}{cd}\$ and \$a * \frac{b}{c} * d\$. But when writing code, on one line, and without parens, 20 m / 5 s is 20 * (m / 5) * s, even though humans might interpret it as (20 m) / (5 s). I used to write kWh / m² . y for annual solar energy, even though it should be kWh / (m² . y). \$\endgroup\$ Sep 14, 2021 at 19:38
  • \$\begingroup\$ Oh gosh darn it, you're so right. Now I feel a little silly :) \$\endgroup\$
    – Peilonrayz
    Sep 14, 2021 at 19:47
  • 2
    \$\begingroup\$ @Peilonrayz: No problem. It feels natural to read it as (20 m) / (5 s), especially when it's written 20m / 5s. It feels weird to separate 5 and s since they're apparently stuck together. There are trick questions on social media based on this ambiguity, e.g. 6 ÷ 2(1+2). The only correct answer is : "use parentheses!". \$\endgroup\$ Sep 14, 2021 at 19:55
3
\$\begingroup\$

One comment about the design: your current list of units

H = ("h", "duration")
MW = ("MW", "power")
MWH = ("MWh", "energy")
EUR = ("Eur", "revenue")
EURMWH = ("Eur/MWh", "price")
DEGC = ("degC", "temperature")

does not separate unit from scale. Energy and power have "mega" baked in, for instance. A more general approach would offer

  • time in units of years, months, weeks, days, hours, minutes, seconds, and SI fractions of seconds;
  • SI multiples for power, energy and temperature; and perhaps
  • Euros or cents (the accepted sub-unit for euros).
\$\endgroup\$
3
  • 1
    \$\begingroup\$ That's a valid point! In this use-case, for energy trade, the units are always the ones mentioned here, and e.g. a GWh is thought of as 1000 MWh, not 1e9 Wh or 3.6e12 J. A more general application would take various prefixes into consideration. BTW: allowing time to be expressed in year, months, quarters or days is extremely tricky, as these do not have a fixed duration. \$\endgroup\$
    – ElRudi
    Sep 15, 2021 at 7:33
  • \$\begingroup\$ Of course time will be tricky, but will sometimes be necessary. You're working in the energy sector: what if you're asked how much energy was consumed over a billing period of one month? \$\endgroup\$
    – Reinderien
    Sep 15, 2021 at 13:28
  • \$\begingroup\$ My point was that, e.g. MWh/month is not a useful unit in and of itself. 720 MWh/month is an average rate of less than, more than, or exactly 1 MW, depending on the month, and therefore only makes sense if the actual month is added to the number. (For this exact use case I've created and added a .duration property to pandas DateatimeIndex and Timestamp objects, which uses the timestamp and its freq argument to find the timedelta in hours. Ah and don't tell me about periodindices. They don't work with DST 😅) \$\endgroup\$
    – ElRudi
    Sep 15, 2021 at 22:06
0
\$\begingroup\$

A small suggestion is to move the relations variable from the function calls to the class and make it static as well. As it will be creating the same variable every time you try to call the __mult__ or __truediv__ this might be overhead.

\$\endgroup\$
2
  • \$\begingroup\$ Thanks for your comment @rishabh-deep-singh. Could you elaborate? My first question proposes using a @staticmethod def relations(): return ((. ,. ,. ),(., ., .)), but that too recreates the variable every time it's called. I suppose I could add @functools.lru_cache, but that's not what you mean, correct? \$\endgroup\$
    – ElRudi
    Sep 13, 2021 at 7:11
  • 1
    \$\begingroup\$ I think @Alex Waygood did a great job in this. \$\endgroup\$ Sep 13, 2021 at 16:01

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.