teepee
teepee

Reputation: 2714

Convert __main__ function from Python script intended for command line into a callable function

I am working with someone's GitHub code that is designed to be called from the command line like so:

> python this_script.py -u <username> -p <password> -i <id_num> ...

This produces an output text file after parsing the inputs in such a manner within this_script.py:

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='get data', add_help=False, usage='this_script.py -u username -p password [options]')
    parser.add_argument('-u', metavar='<STR>', type=str, help='username')
    parser.add_argument('-p', metavar='<STR>', type=str, help='password')
    parser.add_argument('-i', metavar='<STR>', nargs='+', type=str, help='List of IDs')
    ...

I want to use many of these text files within my own set of code, so it would be much more convenient for me to convert this from a command line script to a callable function, with syntax something like this:

def this_script(password, username, *args):
    ...

Is there a simple way to do this, without have to mess around too much with the inner workings of the __main__ block as currently defined?

Upvotes: 1

Views: 1506

Answers (2)

Anna B
Anna B

Reputation: 31

How about:

def do_stuff(args):
    print(args.u)
    print(args.p)


def main(args_list=None):
    parser = argparse.ArgumentParser(description='get data', add_help=False, usage='this_script.py -u username -p password [options]')
    parser.add_argument('-u', metavar='<STR>', type=str, help='username')
    parser.add_argument('-p', metavar='<STR>', type=str, help='password')
    parser.add_argument('-i', metavar='<STR>', nargs='+', type=str, help='List of IDs')
    ...

    if args_list:
        args= parser.parse_args(args_list)
    else:
        args = parser.parse_args()

    do_stuff(args)


if __name__ == "__main__":
    main()

By adding the following line into your package __init__.py

from .__main__ import main

You can call the script as a function from command line, or from another script as long as you import your package, such as

import <package name> as pkg
pkg.main(['-u', 'my_username', '-p', 'my_password'])

Upvotes: 2

wim
wim

Reputation: 363233

Yes, the simple way is to use stdlib runpy.

import runpy

runpy.run_module(mod_name, init_globals=None, run_name=None, alter_sys=False)

And, please make a mental note never to write your own scripts like that. There should only ever be a single line under the conditional, and it should be a call to a function with no arguments, like this:

if __name__ == "__main__":
    main()

Anything more is too much.

Upvotes: 0

Related Questions