Setting a function with a parser command

Question

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.

Answer 1

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))

Answer 2

Here's an example how you might use a base class with your functions as classes inheriting from it to automatically build your commandline (using the __subclasses__() method on the base class) and do the calculation, without writing some ridiculous sequence of if statements.

For more info on subclasses see here in the Python documentation.

import argparse

# base class for all your functions to inherit from    
class mystuff():
        pass

# function classes inherit from mystuff, provide at least a docstring and a calculate() method.            
class fun1(mystuff):
    '''
    Calculate fun1
    '''
    def calculate(self,value):
        return value*2
    
class fun2(mystuff):
    '''
    Calculate fun2
    '''
    def calculate(self,value):
        return value*23

if __name__=="__main__":
    parser = argparse.ArgumentParser(description="Allow calling functions each defined in their own class")

    # find all the classes to be used as commandline arguments
    for fun in mystuff.__subclasses__():
#        print( f"Creating argument for {fun.__name__}" )
        parser.add_argument( f'--{fun.__name__}', default=0, dest='function', action='store_const', const=fun, help=fun.__doc__)
    parser.add_argument('variable', type=int, help="Value to pass to calculation")
    
    args = parser.parse_args()

    # create an instance and call its calculate function passing the variable value
    result = args.function().calculate(args.variable)
    
    print( f"Calling {args.function.__name__} with variable {args.variable} gives result {result}" )

    

Getting the usage:

fun1.py -h

gives:

usage: fun1.py [-h] [--fun1] [--fun2] variable

Allow calling functions each defined in their own class

positional arguments:
  variable    Value to pass to calculation

optional arguments:
  -h, --help  show this help message and exit
  --fun1      Calculate fun1
  --fun2      Calculate fun2

invoking one of the functions:

fun1.py --fun1 235

gives result:

Calling fun1 with variable 235 gives result 470

I'm sure this could be made more sophisticated, perhaps removing the need for a -- before each function name by using subparsers.