ActiveState Code

Recipe 499333: Decorator for main method

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

Discussion

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.

Comments

  1. 1. At 2:05 p.m. on 27 dec 2006, Steven Bethard said:

    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.

Sign in to comment