Question

In a Python script, I would like to control the importing of a module into the main() from a sub function. Is this possible? How?

Why: I have a sub function that handles command line arguments (using argparse) and would like to import a module based on user input. Specifically, I would like to let the user specify the backend for matplotlib, which has to be set before importing matplotlib.pylab. However, I think the question has more general use.

Here is a code snippet:

def main():

    args = handleCommandLine();

    fig, ax = plt.subplots(1)   # ERROR: plt not defined

    # Snip ...


def handleCommandLine():
    p = argparse.ArgumentParser()
    p.add_argument('--backend', '-b', default=None, help='Specify plotting backend')
    args = p.parse_args()

    if args.backend != None:
        matplotlib.use(args.backend)  #Must occur before importing pyplot

    import matplotlib.pyplot as plt   #Must occur after setting backend, if desired
    return args

Answer 1

If you’d like it to behave as though you executed import matplotlib.pyplot as plt at the top of the module, even though you didn’t, use a global:

def handleCommandLine():
    p = argparse.ArgumentParser()
    p.add_argument('--backend', '-b', default=None, help='Specify plotting backend')
    args = p.parse_args()

    if args.backend != None:
        matplotlib.use(args.backend)  #Must occur before importing pyplot

    global plt  #Style choice: Can also be placed at the top of the function
    import matplotlib.pyplot as plt  #Must occur after setting backend
    return args

Otherwise you can pass the included library reference through a function return, in the same way you handle any other variable:

def main():
    plt, args = handleCommandLine()   # CHANGED HERE
    fig, ax = plt.subplots(1)
    # ...    

def handleCommandLine():
    p = argparse.ArgumentParser()
    p.add_argument('--backend', '-b', default=None, help='Specify plotting backend')
    args = p.parse_args()

    if args.backend != None:
        matplotlib.use(args.backend)  #Must occur before importing pyplot

    import matplotlib.pyplot as plt   #Must occur after setting backend, if desired
    return plt, args   # CHANGED HERE