Welcome, guest | Sign In | My Account | Store | Cart

I found an article by Guido a while ago describing a generic main method for most Python programs (http://www.artima.com/weblogs/viewpost.jsp?thread=4829). Since then I found myself copying the same code into most of my programs. This recipe shows how to do the same thing with a decorator that makes the code much cleaner.

Python, 46 lines
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import getopt,sys,traceback
class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg
    def __repr__(self):
        return self.msg

def main_function(getOptString='', numArgs=0):
    def main_decorator(maincode):
        def decorated_main(argv=[__name__]):
            try:
                opts={}
                args=[]
                if len(getOptString) > 0:
                    try:
                        opt_list, args = getopt.getopt(argv[1:], getOptString)
                        opts=dict([(x[0][1:], x[1]) for x in opt_list])
                    except getopt.error, msg:
                        raise Usage(msg)
                else:
                    args=argv[1:]
                if len(args) < numArgs:
                    raise Usage("Not enough arguments")
                return maincode(args, opts) or 0
            except KeyboardInterrupt:
                sys.exit(-1)
            except SystemExit:
                pass
            except Usage, usage:
                sys.stderr.write('%s\n%s\n' % (maincode.__doc__, usage.msg))
            except Exception, err:
                cla, exc, trbk = sys.exc_info()
                import traceback
                sys.stderr.write("Caught Exception:\n%s:%s\n%s" % (cla.__name__, str(exc), ''.join(traceback.format_tb(trbk,5))))
                sys.exit(-1)
        return decorated_main
    return main_decorator

########### Example Usage ################

@main_function('o:', 2)
def main(args, opts):
    """Usage: test.py [-o opt1] arg1 arg2"""

if __name__=="__main__":
    sys.exit(main(sys.argv) or 0)

This allows option and argument parsing to be hidden from the main code in your program. The user simply passes an optional getopt string and the number of required arguments. Then the actual main function get passed a dictionary of options and a list of arguments.

One possible improvement would be to add an optional error handling routine passed to the decorator. As it is all exceptions are caught in the decorator function and printed to the screen.

3 comments

Steven Bethard 14 years, 11 months ago  # | flag

a job for argparse or optparse. If you're really serious about command-line parsing, you should be using argparse (http://argparse.python-hosting.com/) or optparse (from the standard library). Your code could be as simple as:

def main(foo, bars):
    # do something with foo and bars

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--foo')
    parser.add_argument('bars', nargs=2)
    args = parser.parse_args()
    main(**vars(args))

Now if the command line args are "--foo f b1 b2" then main will get called like

main(foo='f', bars=['b1', 'b2'])

Using argparse or optparse means that giving your program better command-line documentation is as simple as adding a help= keyword argument to your add_argument() calls.

Karl Dickman 11 years, 11 months ago  # | flag

I prefer to split the argument parsing into a standalone function. This is returns a tuple (options, arguments), the same way optparse.OptionParser.parse_args() does. I don't have any exception handling, because I like to use optparse. optparse.OptionParser.error(message) halts execution and prints a useful usage message in addition to your custom complaint.

def main_function(parse_arguments=None):
    if parse_arguments is None:
        parse_arguments = lambda arguments: (None, arguments)
    def main_decorator(to_decorate):
        def decorated_main(arguments=None):
            if arguments is None:
                arguments = sys.argv
            options, arguments = parse_arguments(arguments)
            sys.exit(to_decorate(options, arguments))
        return decorated_main
    return main_decorator

@main_function(parse_arguments)
def main(options, arguments):
    pass

def parse_arguments(arguments):
    option_parser = optparse.OptionParser()
    option_parser.add_option("-s", "--spam")
    return option_parser.parse_args(arguments[1:])

if __name__ == "__main__":
    main()
Anjum Naseer 10 years, 10 months ago  # | flag

You may be interested in a little Python module I wrote to make handling of command line arguments even easier (open source and free to use) - http://freshmeat.net/projects/commando