How to typehint dataclass field in function signature

69 views Asked by At

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.

2

There are 2 answers

2
Jared Smith On BEST ANSWER

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

1
jsbueno On

If you want to check for the values and not Enum members, sometimes that simply is not possible at all. Keep in mind that typing is static, not runtime, in Python. When you use typing.Literal, at some point in your source code you have to have actually hardcoded that value, so the type-checking mechanism can trickle it down to the point of your function call.

But if the value is being obtained from an external source, be it a configuration variable, a database, whatever, the typing system can't know about it in any way - it does not run the program.

I think the correct thing to do in this case is to actually work with Enum members - whenever you get your int from an external source, convert it to a proper enum member - and that can be statically checked normally.

To convert a "plain" Python value to an enum member, just call the Enum class with the value: Consts(0) - which will also provide you with a runtime check. If type hinting complains at this step, add marks so that it is ignored at static checking type.