Pythonic and mypy compatible way to map str to Enum

Question

my goal is to have a user input a string that is checked to be valid - by defining a dict of allowed strings - and mapped to another string to be stored in a list. Another function is supposed to map these stored strings to Enums and work with those. Here's the code I have:

from __future__ import annotations

from enum import Enum
from typing import List, Type

ONG_SPS_STATES = {
    "pre": "PreReleased",
    "rel": "Released",
    "ret": "Retired",
    "tgt": "Target"
}

sps_states: List[str] = []

def str_to_state(input: str) -> None:
    if (input in ONG_SPS_STATES.keys()):
        sps_states.append(ONG_SPS_STATES[input])
    else:
        print('Invalid input')

SpsState = Enum('SpsState', list(ONG_SPS_STATES.values()))

def get_sps_state(index: int) -> Type[Enum]:
    if (len(sps_states) > index):
        return SpsState[sps_states[index]]
    raise Exception('Invalid index')

if (__name__ == '__main__'):
    cont = True
    while cont:
        inp = input('State: ')
        if (inp.lower() != 'e'):
            str_to_state(inp)
        else:
            cont = False
    cont = True
    while cont:
        inp = input('Index: ')
        if (inp.lower() != 'e'):
            print(get_sps_state(int(inp)))
        else:
            cont = False

The problem is that mypy complains about line 21 (SpsState = Enum('SpsState', list(ONG_SPS_STATES.values()))) with error: Enum() expects a string, tuple, list or dict literal as the second argument [misc] and line 25 (return SpsState[sps_states[index]]) with error: Incompatible return value type (got "SpsState", expected "Type[Enum]") [return-value]

How can the existing code with the described "ease of use" be rewritten to be pythonic (Python 3.9 if that is important) and compatible with mypy at the same time?

Answer 1

Your enum setup is kind of inside-out and backwards. Don't make a dictionary and don't define an enum via sequence; make a normal Enum subclass. Its values will be your current keys, and its names will be your current values.

Don't surround if predicates with parens in Python.

Don't use a cont loop flag when you can just break.

Suggested

from enum import Enum


class SpsState(Enum):
    PreReleased = 'pre'
    Release = 'rel'
    Retired = 'ret'
    Target = 'tgt'

    @classmethod
    def from_index(cls, idx: int) -> 'SpsState':
        return tuple(cls)[idx]


def input_state_str() -> None:
    while True:
        command = input('State, or "e" to end: ').lower()
        if command == 'e':
            break
        try:
            print(SpsState(command))
        except ValueError as e:
            print('Invalid input:', e)


def input_state_int() -> None:
    while True:
        command = input('Index, or "e" to end: ').lower()
        if command == 'e':
            break
        try:
            print(SpsState.from_index(int(command)))
        except (ValueError, IndexError) as e:
            print('Invalid input:', e)


def main() -> None:
    input_state_str()
    input_state_int()


if __name__ == '__main__':
    main()

State, or "e" to end: 3
Invalid input: '3' is not a valid SpsState
State, or "e" to end: pre
SpsState.PreReleased
State, or "e" to end: e
Index, or "e" to end: 9
Invalid input: tuple index out of range
Index, or "e" to end: 1
SpsState.Release
Index, or "e" to end: e

mypy is OK with it:

Success: no issues found in 1 source file