Generate function signature in python from dataclass

I want to have single file that defines a bunch of constants so that I don't need to hardcode them later on.

For example, I would do this:

@dataclass(frozen=True)
class CONSTANTS:
    STATUS_SUCCESS = 200
    STATUS_ERROR = 400

    SI_UNIT_MASS = "kg"
    SI_UNIT_LENGTH = "m"

Now let's say I have the following function:

def func():
    something = True
    if something:
        return CONSTANTS.STATUS_SUCCESS
    else:
        return CONSTANTS.STATUS_ERROR

I would like to have type hints for func. So I thought I could just add this:

@dataclass(frozen=True)
class TYPES:
    STATUS = Literal[CONSTANTS.STATUS_SUCCESS, CONSTANTS.STATUS_ERROR]
    SI_UNIT = Literal[CONSTANTS.SI_UNIT_MASS, CONSTANTS.SI_UNIT_LENGTH]

and update the functions signature:

def func() -> TYPES.STATUS:

But here I get this error:

Variable not allowed in type expressionPylancereportInvalidTypeForm

The error goes away when I write:

@dataclass(frozen=True)
class TYPES:
    STATUS = Literal[200, 400]
    SI_UNIT = Literal["kg", "m"]

But this has the downside that I might forget to add a new unit or status code to the type. I'd like this to be dynamic so that I don't have to repeat myself.

Is there a standard approach to do this? I also thing I would like to rewrite CONSTANTS in a way were I would like to have

@dataclass(frozen=True)
class CONSTANTS:
    STATUS = STATUS
    SI_UNIT = SI_UNIT

with

@dataclass(frozen=True)
class STATUS:
    SUCCESS = 200
    ERROR = 400

...

So my thought was that I should maybe have OPTIONS and STATUS inherit from a base class that defines something like a TYPE property. But I can't get an example to work that feels natural.

Solution

I can show you why this gives an error, but this does not solve your problem.

First, what is the error? It says "variable not allowed in type expression". Why this makes sense is since in python nothing is protected (or private) all variables can change. For example:

print(func())
# >>> 200

CONSTANTS.STATUS_SUCCESS = "SUCCESS"

print(func())
# >>> SUCCESS

The same can happen for the TYPES.STATUS variable, what happens when the user sets it to TYPES.STATUS = None?

You can either divide the classes and use ENUM where it makes more sense. Or you can look in the @overload decorator, which can be used to define different return types, depending on for example in the input types.

Here is a solution with using enums:

from enum import IntEnum, StrEnum


class STATUS(IntEnum):
    SUCCESS = 200
    ERROR = 400


class UNIT(StrEnum):
    SI_UNIT_MASS = "kg"
    SI_UNIT_LENGTH = "m"


def func() -> STATUS:
    something = True
    if something:
        return STATUS.SUCCESS
    else:
        return STATUS.ERROR


print(func())
# >>> 200

According to the python docs: An Enum is a set of symbolic names bound to unique values. They are similar to global variables, but they offer a more useful repr(), grouping, type-safety, and a few other features.

And in my word I want to say that a dataclass is better suitable as a method how to store data, for example a User class where you have multiple users. Enum classes contain predefined data (in this case SUCCES and ERROR). That is the reason why I suggest the switch to ENUM.

The return value is now a ENUM, but if you use it in a string it will print the value. Therefore print(func()) -> 200 but print(func().__repr__()) -> <STATUS.SUCCESS: 200>.