Search code examples
pythonenumsflags

Pythonic way to implement a selection flag


I work on a Python module where more than 3/4 of my functions offer a choice of 6 options and behave differently accordingly.1

My goal is to make te usage of this argument as simple, obvious and readable as possible, as well as (though less importantly) the implementation of it.

What would be the most pythonic way to implement such a flag?

The way I see it, I have 3 options:

1. Use a string identifier for each component (current solution):
+: Easy and readable usage, no additional imports, classes, implementation details
-: String comparsions, inflexible, user must know the strings (no editor help)

# in module/__init__.py
E_X = 'e_x'
E_Y = 'e_y'
...

# in module/whatever.py:
def do_sg(args, comp):
    if comp == E_X:
        set_some_state_for_ex()
        res = calc_sg_with_ex(args)
    elif comp == E_Y:
        ...
    return res

# usage
from module.whatever import do_sg
res = do_sg(args, 'e_r')

2: Use module-wide constants:
+: literal definition, editor hints
-: More obscure usage, namespace pollution (avoid with more obscure names), no easy way to import all components

# in module/__init__.py
E_X_COMP = 1
E_Y_COMP = 2
...

# in module/whatever.py, implementation v1
from . import *
def do_sg(args, comp):
    if comp == E_X_COMP:
        set_some_state_for_e_x()
        res = calc_sg_with_e_x(args)
    elif comp == E_Y_COMP:
        ...
    return res

# usage ( (1) may be used this way as well)
import module   # or from module import E_X_COMP, E_Y_COMP, ...
from module.whatever import do_sg
res = do_sg(args, module.E_X_COMP)

# or more obscurely:
res = do_sg(args, 0)

3: Use the Enum class:
+: Enclosed literal definition, more obvious usage than constants, better editor hints, type hints
-: More obscure implementation, module dependency, more imports and objects than with string identifiers

# in module/__init__.py
from enum import Enum
class component(Enum):
    e_x = 1
    e_y = 2
    ...

# in module/whatever.py, implementation v1
from . import component
def do_sg(args, comp):
    if comp is component.e_x:
        res = calc_sg_with_e_x(args)
    elif comp is component.e_y:
        ...
    return res

# usage
from module import component
from module.whatever import do_sg
res = do_sg(args, component.e_x)

4: Other options?

Are there any other ways that I didn't consider?


1: To be more specific, the module is for calculating components of electromagnetic fields, and the choice is which component should it consider - which can be x, y, and z components of the E and B field. Vector implementation won't work for various reasons.


Solution

  • Enum was created for exactly this kind of problem. They should also be considered constants, so appropriate naming will help:

    class Component(Enum):
        E_X = 1
        E_Y = 2
    

    And in use;

    from module import Component
    from module.whatever import do_sg
    res = do_sg(args, Component.E_X)