Simple Python argparse wrapper

When creating Python modules I frequently turn them into a Unix-style command-line applications. It makes them easy to demo, test, debug, pipe together, parse the input/output with Unix tools, etc. The argparse module available in Python versions 2.7 and later is great for this sort of thing, but I found myself copying and pasting a lot of code for each new application and thereby breaching the DRY (Don’t Repeat Yourself) principle. I used plac for a while, but then when I went to deploy the application to production I needed to switch it over to argparse. As a result I decided to build my own argparse wrapper:

• Ability to display multi-line version information strings (argparse ignores line feeds), used with the --version option.
• Ability to include positional filename arguments with very little code.
• Ability to allow the last positional filename argument to be replaced with Stdin.
• Ability to call common options from a standard command-argument library.

Examples:

#!/usr/bin/python
__version__ = """argparse wrapper sample code version 0.1a

This is free software.  There is NO warranty; 
not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
"""

USAGE = """cobol2csv.py [OPTIONS] COPYBOOK [DATAFILE]
COPYBOOK - Filename: output from copybook2csv.py
DATAFILE - Filename: COBOL records, fixed-width text
"""

import load

def main(args):
    fields = load.csv_(args.copybook, strip_="right", prune=True)
    data = load.lines(args.datafile, stop_at_line=1)
    if args.verbose:
        print 'In verbose mode...'
    if args.license:
        print 'GPL (GNU Public License)'

if __name__ == '__main__':
    from cmd_line_args import Args
    args = Args(USAGE, __version__)
    args.allow_stdin()
    args.add_files('copybook', 'datafile')
    args.add_options('debug', 'verbose')
    args.parser.add_argument('--license', action='store_true', 
        help='display license information')
    main(args.parse())

Explanation:

• Adds 2 positional filename parameters: copybook & datafile
• The ‘datafile’ parameter can be optionally omitted and Standard Input (Stdin) can be used instead of reading from a file. For example cat file1.txt|./cobol2csv.py layout.csv will work.
• -d, --debug, -v, --verbose are automatically added from a standard library of command-line options
• The argparse object can be accessed directly to support all normal argparse methods & functionallity (i.e. see --license option).

Note: The source code below may be dated, to get the most current version visit Harbingers-Hollow at Github.

import argparse, sys

__all__ = ['Args']

class VersionAction(argparse.Action):
    """Overrides argparse_VersionAction(Action) to allow line feeds within
    version display information.""" 
    def __init__(self, option_strings, version=None, 
         dest=None, default=None, help=None):
         super(VersionAction, self).__init__(option_strings=option_strings,
             dest=dest, default=default, nargs=0, help=help)
         self.version = version

    def __call__(self, parser, namespace, values, option_string=None):
        version = self.version
        if version is None:
            version = parser.version
        print version
        parser.exit()

class Args:
    """argparse wrapper"""
    
    allow_stdin = False
    
    def __init__(self, usage, version):
        self.parser = argparse.ArgumentParser(usage=usage)
        self.parser.version = version
        self.parser.add_argument('-V', '--version', 
            action = VersionAction,
            help='display version information and exit')
    
    def add_files(self, *file_args):
        """Add positional filename argurments.  If self.allow_stdin is set
        Example:
            object.add_filenames('config_file', 'data_file'])
            The 1st filename will be saved in a variable called 'config_file'.
            The 2st filename will be saved in a variable called 'data_file'.
        """
        if self.allow_stdin:
            for file_arg in file_args[:-1]:
                self.parser.add_argument(file_arg, help='filename... %s' % file_arg)
            self.parser.add_argument(file_args[-1], 
                help='filename... %s' % file_args[-1], nargs='?')
        else:
            for file_arg in file_args:
                self.parser.add_argument(file_arg, help='filename... %s' % file_arg)      
        self.file_args = file_args
  
    def add_filelist(self):
        """FUTURE: ability to add a list of files to be processed.  Similar to
        python filelist module, but with ability to include other arguments.
        Support for wildcards."""
        pass
        
    def add_options(self, *options):
        """Add from a standard library of pre-defined command-line arguments"""
        for option in options:
            option = option.lower()
            if option == 'debug':
                self.parser.add_argument('-d', '--debug', action='store_true',
                    help='Turn on debug mode.')
            elif option == 'debug_level':
                self.parser.add_argument('-d', '--debug', type=int,
                    help='Set debug level 1-10.')
            elif option == 'verbose':
                self.parser.add_argument('-v', '--verbose', action='store_true',
                    help='Turn on verbose mode.')
            elif option == 'quiet':
                self.parser.add_argument('-q', '--quiet', action='store_true',
                    help='Suppress all output to terminal.')
    
    def allow_stdin(self):
        self.allow_stdin = True

    def parse(self):
        """Parse args & use sys.stdin if applicable
        Sets all file arguments to a file read object"""
        args = self.parser.parse_args()
        if self.file_args:
            if self.allow_stdin:
                if not sys.stdin.isatty():
                    setattr(args, self.file_args[-1], sys.stdin)
                else:
                    self.allow_stdin = False
            last_arg_idx = len(self.file_args) - self.allow_stdin
            for file_arg in self.file_args[:last_arg_idx]:
                try:
                    file_ = open(getattr(args, file_arg))
                except IOError, error_msg:
                    sys.stderr.write('ERROR loading file "%s".\n%s\n' % 
                        (file_arg, error_msg))
                    sys.exit(1)
                setattr(args, file_arg, file_)        
        return args   

Tags: Command-line, Python