Search code examples
pythonpython-typingpython-dataclasses

How to typehint dataclass field in function signature


I've a frozen dataclass with some constants and function, which accept any constant from this class as it's argument.
How can I typehint this mechanic? I need to tell user that function waits for any field's value from this specific dataclass.
It should looks like this:

from dataclasses import dataclass


@dataclass(frozen=True)
class Consts:
    const_0: int = 0
    const_1: int = 1
    const_2: int = 2
    const_3: int = 3


def foo(param: Consts.field):
    return param

UPD:
According to @JaredSmith hint I've tried to use Enum. It look like much correct. But the problem still exists.
I've tried to use typing.Literal, like this:

def foo(param: Literal[Consts.const_0, Consts.const_1]):
    return param

but it will not give use a correct typehint. In case foo(Consts) we will get this, not pretty obvious, warning: Expected type 'Consts', got 'Type[Consts]' instead instead of something like that: Expected type Consts.value got Consts instead

So, the main question, after updates, is: How to aggregate constants in logical groups in the code to simplify their usage (dataclass or Enum) and how to typehint the corresponding solution.

UPD: Great thanks for all comments. They showed me lot of interesting facts. As an answer for current question I choose IntEnum variant by @JaredSmith. Also, thanks @chepner and @InSync for the deep explanation.


Solution

  • You want to use an enum here rather than a dataclass:

    from enum import IntEnum
    
    class MyEnum(IntEnum):
        A = 1
        B = 2
    
    def foo(x: MyEnum):
        if (x == 1):
            return '1 given'
        else:
            return 'whatever'
    
    foo(MyEnum.A) # Ok
    foo(1)        # Error 
    foo('pizza')  # Error
    

    Tested against pyright and mypy type-checkers (I don't use PyCharm, not sure if it's running one of those under the hood or something custom to JetBrains).