Python import based on input parameters

Question

I have a seemingly simple Python code design problem that I have not found any elegant solution to. I have a set of modules (two in the example below) that define functions of the same name that compute the same thing but using different algorithms:

algorithm_a.py:

def compute_1(x):
   # do some computation
   return (x+1.)**2

def compute_2(x):
   # do some other computation

# etc

algorithm_b.py:

def compute_1(x):
   # do same computation in another way
   return x*x + 2.0*x + 1.0

def compute_2(x):
   # do some other computation in another way

# etc

Both modules have approximately 10 functions, but the number might grow in the future.

In a third module, I have functions defined that are not aware of the details of how the computation is performed, i.e. they only care about getting the correct answer. The interface to the user is provided by a class. All in all, the third module looks similar to:

simulation.py:

import XXX as algo #????

def do_computation_1(x):

   v1 = algo.compute_1(x)
   v2 = algo.compute_2(x)

   # Do some computations using v1, v2 
   return result

def do_computation_2(x):
   return algo.compute_2(x)

# etc etc

class Simulation(object):

  def __init__(self, method):
     # Do some magic here so that
     # if method == "A":
     #   algo = algorithm_a
     # if method == "B"
     #   algo = algorithm_b

  def run(self, x):
     do_computation_1(x)
     do_computation_2(x)

How can I get the correct module loaded and applied in do_computation() to depend on the method parameter supplied to the class?

Note that the do_computation functions need to remain outside of the Simulation class.

CONCLUSION: The comment by @BrenBarn below summarises well the available options. Thanks all for the great help!

Answer 1

With the way your code is structured, you can't have the import depend on a parameter passed to the class. Python files are executed top-to-bottom. Your import occurs before the class definition, so by the time the class is defined, the import has already occurred. (And you won't be passing method in until you instantiate the class, which will be even later.)

If it's okay to import both modules and you just want to use the specified one, you can do almost literally what you wrote in the comments:

import algorithm_a
import algorithm_b

class Simulation(object):

    def __init__(self, method):
        if method == "A":
            self.algo = algorithm_a
        if method == "B"
            self.algo = algorithm_b

    def do_computation(self, x):
        return self.algo.compute(x)

    def run(self, x):
        self.do_computation(x)

(I have here made do_computation a method of the class. It doesn't make much sense to have it as a separate function if its behavior is determined by a parameter passed to the class.)

If the actual import of the modules may be slow, you could conditionally import one module or the ohter as shown in Reblochon's answer. However, to do this you must put the imports inside the class. If you are going to specify the algorithm via something like Simulation("A"), there's no way to determine which import to do at the time you do import simulation, because you haven't yet specified the algorithm at that time; you would have to wait until you actually call Simulation("A"). (If the import is slow, this will cause a slowdown at that point, which may not be desirable.)

Edit: if you really need to have do_computation be a global function because of Numba, you can work around it by setting a global variable. Change the __init__ to:

def __init__(self, method):
    global algo
    if method == "A":
        algo = algorithm_a
    if method == "B"
        algo = algorithm_b

And make a global do_computation function like this:

def do_computation(x):
    return algo.compute(x)

This will be more confusing because every time you create a new Simulation it will change the global behavior, even for previously-created simulations. But if you aren't creating multiple simulations with different algorithms at the same time, it should be okay.

Answer 2

The better way to do this is to actually save (or pass) the function you want to use.. E.g.

import algorithm_a
import algorithm_b


class Simulation(object):

  def __init__(self, method):
     # Do some magic here so that
     if method == "A":
       self.compute_func = algorithm_a.compute
     if method == "B"
       self.compute_func = algorithm_b.compute

  def run(self, x):
     self.compute_func(x)

If you really must have your external def do_computation(x) function you can pass the algorithm you want to use as an argument

def do_computation(x, compute):
   return compute(x)

class Simulation(object):

  def __init__(self, method):
     # Do some magic here so that
     if method == "A":
       self.compute_func = algorithm_a.compute
     if method == "B"
       self.compute_func = algorithm_b.compute

  def run(self, x):
     do_computation(x, self.compute_func)