Search code examples
pythoninheritanceenumstype-hinting

Using Enum values as type variables, without using Literal


I'm trying to represent physical dimensions (length, time, temperature, ...), and cannot find a nice way to do so, that is compatible with type hinting and generics.

Ideally, I want to be able to define an Enum whose names are types themselves (a metaenum ?):

from enum import Enum

class Dim(Enum):
    TIME = "t"
    MASS = "m"

I can type hint dimensions (dim: Dim) but cannot do things like

from typing import Generic, TypeVar

T = TypeVar("T", bound=Dim)  # only accepts `Dim`

class PhysicalQuantity(Generic[T]):
    pass

class Container:
    some_time: PhysicalQuantity[Dim.TIME]  # doesn't work

because these are values.

Is there a construct as simple as Enum, but to make types instead of values ?

Reasons why I want to keep Enum:

  • very easy to define
  • very easy to associate to a value (str)
  • Ability to sort of "think of Dim as the type, and Dim.TIME as a subtype"

There are functional solutions, however I'm asking this to get a "best way" more than a "working way". Here's what I found:

  1. The simplest solution is to do use Literal: SomeGenericType[Literal[Dim.TIME]], but this is both annoying to write each time and counter-intuitive for people who expect Dim.TIME to behave as a type.

  2. Switching to classes, the most intuitive idea:

    class Dimension: 
        pass
    
    class TIME(Dimension): 
        pass
    

    doesn't work, because I want type(TIME) to be Dim, to reproduce Enum behavior

  3. That leads to using a metaclass:

    class Dimension(type):
        # ... complete __init__ and __new__ to get TIME.symbol = "t"
    
    class TIME(metaclass=Dimension, symbol="t"): 
        pass
    

    This works, but I lose the ability to do Dim.TIME, to get Dim.TIME from Dim('t'), ...


Solution

  • Is there a construct as simple as Enum, but to make types instead of values?

    Yes, the metaclass. A metaclass makes types. It is simple in terms of usage i.e. creation of new types, but you do need to put in some more work to set it up properly.

    Semantically, you could think of the Dimension is a type and Time, Distance etc. as instances of it. In other words the type of the Time class is Dimension. This seems to reflect your view since you said:

    I want type(Time) to be Dim

    Now a Quantity could be considered the abstract base class of type Dimension. Something without a symbol.

    Time would inherit from Quantity (thus also being of type Dimension). No generics so needed so far.

    Now you can define a Container that is generic in terms of the type(s) of quantity (i.e. instances of Dimension) it holds.

    The metaclass and base class could look like this:

    from __future__ import annotations
    from typing import Any, ClassVar, TypeVar, overload
    
    T = TypeVar("T", bound=type)
    
    
    class Dimension(type):
        _types_registered: ClassVar[dict[str, Dimension]] = {}
    
        @overload
        def __new__(mcs, o: object, /) -> Dimension: ...
    
        @overload
        def __new__(
            mcs: type[T],
            name: str,
            bases: tuple[type, ...],
            namespace: dict[str, Any],
            /,
            **kwargs: Any,
        ) -> T: ...
    
        def __new__(
            mcs,
            name: Any,
            bases: Any = None,
            namespace: Any = None,
            /,
            **kwargs: Any,
        ) -> type:
            if bases is None and namespace is None:
                return mcs._types_registered[name]
            symbol = kwargs.pop("symbol", None)
            dim = super().__new__(mcs, name, bases, namespace, **kwargs)
            if symbol is not None:
                mcs._types_registered[symbol] = dim
            return dim
    
    
    class Quantity(metaclass=Dimension):  # abstract base (no symbol)
        pass
    

    And to create new Dimension classes you just inherit from Quantity:

    from typing import Generic, TypeVar, reveal_type
    
    # ... import Quantity
    
    
    class Time(Quantity, symbol="t"):
        pass
    
    
    DimTime = Dimension("t")
    print(DimTime)        # <class '__main__.Time'>
    print(type(Time))     # <class '__main__.Dimension'>
    reveal_type(DimTime)  # mypy note: Revealed type is "Dimension"
    
    
    Q = TypeVar("Q", bound=Quantity)
    
    
    class Container(Generic[Q]):
        """generic container for `Dimension` instances (i.e. quantities)"""
        some_quantity: Q
    

    I realize this completely bypasses your Enum question, but since you even phrased the question as an XY Problem yourself by explaining what your actual intent was, I thought I'd give it a go and suggest a different approach.