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

• user2606240: Ah, I was not aware that you could construct an Enum instance from its value like that! Regarding Don't surround if predicates with parens in Python.: Why is that? I find it way more readable that way - which might just be because I am so used to parentheses after the if from C...
• Reinderien: Whether it's more legible or not (it's not), it's not Pythonic. Linters will correctly tell you to drop the redundant parens, and you should be both using and listening to a linter.