Setting a function with a parser command

When I launch my code I would like to choose a function (from a set of functions) which will be used. Unfortunately I have many loops and the code is very expensive so the following pseudocode is is highly deprecated:

import argparse
import mystuff


code body in which my_variable gets created

if args.my_index == 1:
  def my_func(my_variable):
    return my_variable + 1

if args.my_index == 2:
  def my_func(my_variable):
    return my_variable**2 +1

having used the following command:

$ python3 my_code.py --index 1

I was thinking about promoting the function to an external class module, maybe using the properties of class initialization.

Solution

You can register your functions inside a container like a tuple. Then your can retrieve them by index. the .index attribute of your ArgumentParser object is going to be 1 more than the tuple indices:

import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--index', type=int)
args = parser.parse_args()

my_variable = 10
funcs = (lambda x: x + 1, lambda x: x ** 2 + 1, lambda x: x ** 3 + 1)

if args.index is not None:
    print(funcs[args.index - 1](my_variable))

This way when you execute your script using python3 my_code.py --index 1, the .index is 1, so you need to get the first item of the tuple which is args.index - 1.

output: 11

If by any chance your functions follow a specific pattern(like my_variable ** n + 1 here) you can define a generic function that handles it without registering all the functions:

import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--index', type=int)
args = parser.parse_args()

my_variable = 10

def func(x):
    return my_variable ** x + 1

if args.index is not None:
    print(func(args.index))