moulinette/parse_args

#!/usr/bin/env python
# -*- coding: utf-8 -*-

__credits__ = """
    Copyright (C) 2012 YunoHost

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published
    by the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program; if not, see http://www.gnu.org/licenses
    """
__author__  = 'Kload <kload@kload.fr>'
__version__ = '2.0 beta1'

import os
import sys
import argparse
import gettext
try:
    import yaml
except ImportError:
    sys.stderr.write('Error: Yunohost CLI Require yaml lib\n')
    sys.stderr.write('apt-get install python-yaml\n')
    sys.exit(1)
import yaml
import json
if not __debug__:
    import traceback

gettext.install('YunoHost')

try:
    from yunohost import YunoHostError, YunoHostLDAP, str_to_func, colorize, pretty_print_dict, display_error, connect_services, disconnect_services
except ImportError:
    sys.stderr.write('Error: Yunohost CLI Require YunoHost lib\n')
    sys.exit(1)


def parse_dict(action_map):
    """
    Turn action dictionnary to parser, subparsers and arguments

    Keyword arguments:
        action_map -- Multi-level dictionnary of categories/actions/arguments list

    Returns:
        Namespace of args

    """
    # Intialize parsers
    parsers = subparsers_category = subparsers_action = {}
    parsers['general'] = argparse.ArgumentParser()
    subparsers = parsers['general'].add_subparsers()

    # Add general arguments
    for arg_name, arg_params in action_map['general_arguments'].items():
        if arg_params['full']:
            arg_fullname = arg_params['full']
            del arg_params['full']
            parsers['general'].add_argument(arg_name, arg_fullname, **arg_params)
        else:
            parsers['general'].add_argument(arg_name, **arg_params)

    del action_map['general_arguments']

    # Split categories into subparsers
    for category, category_params in action_map.items():
        if 'category_help' not in category_params: category_params['category_help'] = ''
        subparsers_category[category] = subparsers.add_parser(category, help=category_params['category_help'])
        subparsers_action[category] = subparsers_category[category].add_subparsers()
        # Split actions
        if 'actions' in category_params:
            for action, action_params in category_params['actions'].items():
                if 'action_help' not in action_params: action_params['action_help'] = ''
                parsers[category + '_' + action] = subparsers_action[category].add_parser(action, help=action_params['action_help'])
                # Set the action s related function
                parsers[category + '_' + action].set_defaults(
                    func=str_to_func('yunohost_' + category + '.'
                                     + category + '_' + action))
                # Add arguments
                if 'arguments' in action_params:
                    for arg_name, arg_params in action_params['arguments'].items():
                        if 'full' in arg_params:
                            arg_fullname = arg_params['full']
                            del arg_params['full']
                            parsers[category + '_' + action].add_argument(arg_name, arg_fullname, **arg_params)
                        else:
                            parsers[category + '_' + action].add_argument(arg_name, **arg_params)

    return parsers['general'].parse_args()


def main():
    """
    Main instructions

    Parse the action_dict and execute the action-specific function,
    then print json or pretty result if executed in a tty :)

    Returns:
        int -- 0 or error code

    """
    with open('action_map.yml') as f:
        action_map = yaml.load(f)

    args = parse_dict(action_map)
    connections = connect_services(action_map)

    try:
        if connections:
            result = args.func(vars(args), connections)
        else:
            result = args.func(vars(args))
    except TypeError, error:
        if not __debug__ :
            traceback.print_exc()
        print(_("Not (yet) implemented function"))
        return 1
    except YunoHostError, error:
        display_error(error)
        return error.code
    else: 
        if result is None:
            pass
        elif os.isatty(1):
            pretty_print_dict(result)
        else:
            print(json.dumps(result))
    finally:
        disconnect_services(connections)

    return 0

if __name__ == '__main__':
    sys.exit(main())