Forbid enum entry as a value for function argument with python typing

Question

I have an enum of commands, where all of them have some value except SERVER_CONFIRMATION:

class ServerCommand(StrEnum):
    SERVER_CONFIRMATION = ""
    SERVER_MOVE = "102 MOVE"
    SERVER_TURN_LEFT = "103 TURN LEFT"
    SERVER_TURN_RIGHT = "104 TURN RIGHT"
    SERVER_PICK_UP = "105 GET MESSAGE"
    SERVER_LOGOUT = "106 LOGOUT"
    SERVER_KEY_REQUEST = "107 KEY REQUEST"
    SERVER_OK = "200 OK"
    SERVER_LOGIN_FAILED = "300 LOGIN FAILED"
    SERVER_SYNTAX_ERROR = "301 SYNTAX ERROR"
    SERVER_LOGIC_ERROR = "302 LOGIC ERROR"
    SERVER_KEY_OUT_OF_RANGE_ERROR = "303 KEY OUT OF RANGE"

Then have a class which has to create some text based on the enum member passed to it. For SERVER_CONFIRMATION I need an additional int parameter, while for all other enum members I don't. I tried to write overloads with typing.overload like this:

class CommandCreator(ABC):
    @overload
    @abstractmethod
    def create_message(
        self, cmd: Literal[ServerCommand.SERVER_CONFIRMATION], confirmation_number: int
    ) -> bytes:
        pass

    @overload
    @abstractmethod
    def create_message(self, cmd: ServerCommand) -> bytes:
        pass

    @abstractmethod
    def create_message(self, cmd: ServerCommand, confirmation_number: int | None = None) -> bytes:
        pass

But this still allows for calls without an int parameter like so:

command_creator.create_message(ServerCommand.SERVER_CONFIRMATION)

Of course it is possible to write some monstrosity like this, but it does not seem like a good solution:

ServerCommandWithoutArgument = Literal[
    ServerCommand.SERVER_MOVE,
    ServerCommand.SERVER_TURN_LEFT,
    ServerCommand.SERVER_TURN_RIGHT,
    ServerCommand.SERVER_PICK_UP,
    ServerCommand.SERVER_LOGOUT,
    ServerCommand.SERVER_KEY_REQUEST,
    ServerCommand.SERVER_OK,
    ServerCommand.SERVER_LOGIN_FAILED,
    ServerCommand.SERVER_SYNTAX_ERROR,
    ServerCommand.SERVER_LOGIC_ERROR,
    ServerCommand.SERVER_KEY_OUT_OF_RANGE_ERROR,
]

class CommandCreator(ABC):
    @overload
    @abstractmethod
    def create_message(
        self, cmd: Literal[ServerCommand.SERVER_CONFIRMATION], confirmation_number: int
    ) -> bytes:
        pass

    @overload
    @abstractmethod
    def create_message(self, cmd: ServerCommandWithoutArgument) -> bytes:
        pass

    @abstractmethod
    def create_message(self, cmd: ServerCommand, confirmation_number: int | None = None) -> bytes:
        pass

Is there is a better way to solve this?

Daniil Fajnberg · Answer 1 · 2023-05-20T09:34:38.200

The problem is that you are misusing the Literal annotation by passing your enum member to it. As should be evident from the name, Literal requires you to provide literal values, not variables. To quote from the documentation, it

~~can be used to indicate to type checkers that the corresponding variable or function parameter has a value equivalent to the provided literal [...].~~

~~(Emphasis mine)~~

UPDATE: It turns out PEP 586 explicitly defines support for Enum members as arguments to Literal in addition to actual literals. So your use of Literal is perfectly valid. Thanks to @SUTerliakov for pointing this out.

You have a few options to achieve what you want. One way is to remove the empty string from your enum and instead use Literal ~~correctly~~ in the first overload:

from enum import StrEnum
from typing import Literal, overload


class ServerCommand(StrEnum):
    SERVER_MOVE = "102 MOVE"
    SERVER_TURN_LEFT = "103 TURN LEFT"
    ...


class CommandCreator:
    @overload
    def create_message(self, cmd: Literal[""], confirmation_number: int) -> bytes:
        pass

    @overload
    def create_message(self, cmd: ServerCommand) -> bytes:
        pass

    def create_message(
        self,
        cmd: ServerCommand | Literal[""],
        confirmation_number: int | None = None,
    ) -> bytes:
        return b"foo"

Another option is to split the enum into two. Something like this:

from enum import StrEnum
from typing import overload


class ServerConfirmation(StrEnum):
    CMD = ""


class ServerCommand(StrEnum):
    SERVER_MOVE = "102 MOVE"
    SERVER_TURN_LEFT = "103 TURN LEFT"
    ...


class CommandCreator:
    @overload
    def create_message(self, cmd: ServerConfirmation, confirmation_number: int) -> bytes:
        pass

    @overload
    def create_message(self, cmd: ServerCommand) -> bytes:
        pass

    def create_message(
        self,
        cmd: ServerCommand | ServerConfirmation,
        confirmation_number: int | None = None,
    ) -> bytes:
        return b"foo"

Unfortunately the built-in enums are not extensible. But there are libraries out there that would allow you to. That may offer you another solution.

You could also just not use an enum, but use only actual literals for annotations instead.

And lasty, you could also just not write an overload for this at all, relying instead on good old runtime errors, when the special enum member is passed without an accompanying int.

Why shouldn't I use `Literal` with enum entries like this? How would I do a function overload depending on enum entry if I used `auto()` for example and didn't know/care about values? — Hryhorii Biloshenko, Apr 14 '23 at 16:21
@HryhoriiBiloshenko You don't. I even quoted the docs in my answer to drive home the explanation. You seem to have trouble with the meaning of the word **literal**. `"foo"` is a literal, so is `3.14` and `[1, 10, 100]`, but when you do `x = "foo"` and then use `x` somewhere that `x` is not a literal, that is a variable with a value. I don't know how else to explain this you. `typing.Literal` expects **literals** not variables. The reason is that a static type checker does not execute your code, it just reads it. — Daniil Fajnberg, Apr 14 '23 at 16:26
Then is there anything what I can use instead to allow such overloading? — Hryhorii Biloshenko, Apr 14 '23 at 16:32
@HryhoriiBiloshenko I assume none of the five workarounds I suggested in my answer work for you? — Daniil Fajnberg, Apr 14 '23 at 16:33
From your comment I've found out that I use enum values inappropriately and just enums are not enough, so I guess I'll refactor my code to use classes with enum fields and "overload" constructor for some child classes to take those extra params. Thank you for help. — Hryhorii Biloshenko, Apr 14 '23 at 17:01
The very beginning of this answer is **wrong**. Using enum members as `Literal` items is explicitly allowed by [PEP586](https://peps.python.org/pep-0586/#legal-parameters-for-literal-at-type-check-time) and fully supported by `mypy`, so building a `Literal` of some Enum members subset is certainly not a `Literal` misuse. (e.g. `Literal[ServerCommand.SERVER_CONFIRMATION]` is perfectly valid) — STerliakov, May 20 '23 at 00:37
@SUTerliakov Wow, now I feel like an idiot. I could have just looked this up myself. I did not expect this at all and it still seems strange to make an exception for something that is not actually a literal. But I can see the idea behind this decision. Anyway, thank you for pointing this out! I fixed my post. — Daniil Fajnberg, May 20 '23 at 09:37
I'd go with your last solution anyway, because it shows the clear distinction between variants with and without arguments - they do not belong to the same enum semantically. — STerliakov, May 20 '23 at 10:49

Forbid enum entry as a value for function argument with python typing

1 Answers1