python-sclapp 0.5.3-3 / usr / share / pyshared / sclapp / main.py

# Copyright (c) 2005-2007 Forest Bond.
# This file is part of the sclapp software package.
# 
# sclapp is free software; you can redistribute it and/or modify it under the
# terms of the GNU General Public License version 2 as published by the Free
# Software Foundation.
# 
# A copy of the license has been included in the COPYING file.

'''sclapp's mainWrapper()'''

import os, sys, logging, errno, signal, locale

from sclapp.exceptions import CriticalError, UsageError
from sclapp.signals import enableSignalHandling, disableSignalHandling, \
  getExitSignals
from sclapp.error_output import printDebug, printCritical, setErrorOutputLevel
from sclapp.protected_output import protectOutput, unprotectOutput
import sclapp.daemonize
from sclapp.stdio_encoding import enableStdioEncoding

def mainWrapper(
  realmain,
  name = None,
  author = u'the author',
  version = None,
  doc = None,
  handle_signals = True,
  exit_signals = None,
  notify_signals = None,
  default_signals = None,
  ignore_signals = None,
  protect_output = True,
  daemonize = False,
  bug_message = (
    u'${traceback}\n'
    u'Something bad happened, and is most likely a bug.\n'
    u'Please file a bug report to ${author}.\n'
    u'Include the error message(s) printed here.'
  ),
  version_message = u'${name} version ${version}',
  ignore_unrecognized_options = True,
  error_output_level = None, announce_signals = True,
  decode_stdin = True, encode_stdout = True, encode_stderr = True,
  decode_argv = True,
):
    '''
    Some of the optional parameters mentioned above cause sclapp to parse argv
    and respond appropriately:

    * If version_message is not None, sclapp will respond to the -v command
      line switch by printing the version_message to stdout.
    * If doc is not None, sclapp will respond to the -h command line switch
      by printing doc to stdout.

    Both bug_message and version_message are parsed for substitution strings
    prior to printing.  Specifically, sclapp uses the Template class of the
    standard library's string module to make the following substitutions:

    ${name} is replaced by the name parameter
    ${author} is replaced by the author parameter
    ${version} is replaced by the version parameter
    ${doc} is replaced by the doc parameters

    The doc parameter is parsed for all other substitutions before it is itself
    substituted.  Thus, callers should feel free to add ${name}, ${author}, and
    ${version} to the doc parameter, and they will be replaced appropriately
    before printing.

    Be default, sclapp installs signal handlers that are easily configurable
    using module functions.  Setting handle_signals to False will disable this
    behavior.

    sclapp can also be used to write simple daemons.  To cause the program to
    daemonize before launching the wrapped main() function, set daemonize to
    True.

    bug_message, if not None, will be printed in the event that an unhandled
    exception is caught by the main wrapper.  The default message simply
    informs the user that a likely bug has been encountered, and asks them to
    file a bug report to the author.

    If ignore_unrecognized_options is False, sclapp will complain about options
    it isn't expecting.  If your program does not explicitly handle command
    line options, you probably want this.  Otherwise, you definately don't want
    this, as any of your options will cause a UsageError to be reported before
    you have a chance to handle them.

    decode_stdin, encode_stdout, and encode_stderr determine whether or not
    sclapp enables on-the-fly stdio en/decoding.  They default to True.

    If decode_argv is True, members of the argv argument passed to the main
    function will be decoded to Unicode.  It defaults to True.

    If sclapp finds that it can't do anything useful with -h/--help or
    -v/--version (if version, doc are None, or contain substitutions sclapp
    can't fulfill), sclapp does nothing with command line options.

    The new main() function will return the wrapped function's return value,
    unless an exception is raised (including exceptions resulting from an exit
    signal being received), in which case the return value is an integer
    intended to act as the program's termination status.
    '''

    # functools was added in Python 2.5
    try:
        from functools import wraps
    except ImportError:
        from sclapp.legacy_support import wraps

    # Template was added in Python 2.4
    try:
        from string import Template
    except ImportError:
        from sclapp.legacy_support import Template

    try:
        from traceback import format_exc
    except ImportError:
        from sclapp.legacy_support import format_exc

    encoding = locale.getpreferredencoding()

    argv = sys.argv
    if decode_argv:
        argv = [unicode(a, encoding) for a in argv]

    if name is None:
        name = os.path.basename(argv[0])

    def _sub(message):
        mapping = { }
        if name is not None:
            mapping['name'] = name
        if author is not None:
            mapping['author'] = author
        if version is not None:
            mapping['version'] = version
        mapping['traceback'] = format_exc()
        return Template(message).safe_substitute(mapping)

    def _handleArgv(argv):
        for arg in argv:
            if arg.startswith('-') and (
              arg not in ('-h', '-v', '--help', '--version')):
                if not ignore_unrecognized_options:
                    raise UsageError, u'unrecognized option: %s' % arg
        if (version_message is not None) and (version is not None):
            try:
                version_message_out = _sub(version_message)
            except (KeyError, ValueError), e:
                pass
            else:
                if ('-v' in argv) or ('--version' in argv):
                    print version_message_out
                    raise CriticalError, (0, None)
        if doc is not None:
            doc_out = _sub(doc)
            if doc_out is not None:
                if ('-h' in argv) or ('--help' in argv):
                    print doc_out
                    raise CriticalError, (0, None)

    #@wraps(realmain)
    def main(argv = argv):
        try:
            enableStdioEncoding(
              decode_stdin = decode_stdin,
              encode_stdout = encode_stdout,
              encode_stderr = encode_stderr
            )
            try:
                if error_output_level is not None:
                    setErrorOutputLevel(error_output_level)
                if os.name == 'posix' and handle_signals:
                    enableSignalHandling(
                      exit_signals = exit_signals,
                      notify_signals = notify_signals,
                      default_signals = default_signals,
                      ignore_signals = ignore_signals,
                      announce_signals = announce_signals
                    )

                    if (
                      not protect_output
                    ) and (
                      hasattr(signal, 'SIGPIPE')
                    ) and (
                      signal.SIGPIPE in getExitSignals()
                    ):
                        raise AssertionError, (
                          'Protected output must be enabled to reliably '
                          'handle SIGPIPE as an exit signal.'
                        )

                if daemonize:
                    sclapp.daemonize.daemonize()
                if protect_output:
                    protectOutput()
                _handleArgv(argv)
                status = realmain(argv)
            except KeyboardInterrupt, e:
                raise CriticalError, (0, unicode(e))
            except UsageError, e:
                if doc is not None:
                    doc_out = _sub(doc)
                    print doc_out
                raise
        except SystemExit, e:
            raise
        except CriticalError, e:
            if e.code != 0:
                printCritical(unicode(e))
            status = e.code
        except:
            if bug_message is not None:
                msg = _sub(bug_message)
            else:
                msg = format_exc()
            try:
                printCritical(msg)
            except:
                print msg
                raise
            status = 255
        if protect_output:
            unprotectOutput()
        logging.shutdown()
        if status is not None and status != 0:
            printDebug(u'exiting with status %i' % int(status))
        return status

    main = wraps(realmain)(main)

    return main

def main_function(*args, **kwargs):
    '''Decorator for main functions.  Both of these variations should work:

    @main_function
    def main(argv):
        return 0

    @main_function(name = 'myprog')
    def main(argv):
        return 0
    '''
    if (len(args) == 1) and (len(kwargs) == 0):
        return mainWrapper(args[0])
    elif (len(args) > 1):
        raise ValueError, (
          'main_function accepts either a signle positional arg or zero or '
          'more keyword args'
        )
    def decorator(fn):
        return mainWrapper(fn, **kwargs)
    return decorator

def makeSubCommandMain(cmds, initialize = None, default = None):
    def main(argv):
        if len(argv) < 2:
            if default is None:
                raise UsageError, (
                  u'expected: one of %s' % ', '.join(cmds.keys()))
            else:
                cmd = default
        else:
            cmd = argv[1]

        try:
            cmd_fn = cmds[cmd]
        except KeyError:
            raise UsageError, u'no such sub-command: %s' % cmd

        if initialize is not None:
            initialize(argv)
        return cmd_fn(argv)

    return main