CODEWITHSUNDEEP
pythonenumsflags
Neinstein
Neinstein

Reputation: 1033

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.

Upvotes: 2

Views: 916

Answers (2)

Ethan Furman
Ethan Furman

Reputation: 69288

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)

Upvotes: 1

ghchoi
ghchoi

Reputation: 5166

It is interesting because I like to talk about code style. I do not know what you exactly want to do but it seems like implementing switch in a neat way. If I were you, I would not use any (global) constants. They always make code less readable.

One way to make your code more readable, I think you can define if-else blocks as functions and directly use functions instead of branching with outer scope constants. For example:

# in module/whatever.py, implementation v1
from . import *
class do_sg:
    @classmethod
    def e_x(args):
        set_some_state_for_e_x()
        res = calc_sg_with_e_x(args)
        return res

    @classmethod
    def e_y(args):
        ...
        return res

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

Upvotes: 0

Related Questions