nilmdb/nilmdb/cmdline/cmdline.py

"""Command line client functionality"""

import nilmdb.client

from nilmdb.utils.printf import *
import datetime_tz
import nilmdb.utils.time

import sys
import os
import argparse
from argparse import ArgumentDefaultsHelpFormatter as def_form
import signal

try: # pragma: no cover
    import argcomplete
except ImportError: # pragma: no cover
    argcomplete = None

# Valid subcommands.  Defined in separate files just to break
# things up -- they're still called with Cmdline as self.
subcommands = [ "help", "info", "create", "rename", "list", "intervals",
                "metadata", "insert", "extract", "remove", "destroy" ]

# Import the subcommand modules
subcmd_mods = {}
for cmd in subcommands:
    subcmd_mods[cmd] = __import__("nilmdb.cmdline." + cmd, fromlist = [ cmd ])

class JimArgumentParser(argparse.ArgumentParser):
    def parse_args(self, args=None, namespace=None):
        # Look for --version anywhere and change it to just "nilmtool
        # --version".  This makes "nilmtool cmd --version" work, which
        # is needed by help2man.
        if "--version" in (args or sys.argv[1:]):
            args = [ "--version" ]
        return argparse.ArgumentParser.parse_args(self, args, namespace)

    def error(self, message):
        self.print_usage(sys.stderr)
        self.exit(2, sprintf("error: %s\n", message))

class Complete(object): # pragma: no cover
    # Completion helpers, for using argcomplete (see
    # extras/nilmtool-bash-completion.sh)
    def escape(self, s):
        quote_chars = [ "\\", "\"", "'", " " ]
        for char in quote_chars:
            s = s.replace(char, "\\" + char)
        return s

    def none(self, prefix, parsed_args, **kwargs):
        return []
    rate = none
    time = none
    url = none

    def path(self, prefix, parsed_args, **kwargs):
        client = nilmdb.client.Client(parsed_args.url)
        return ( self.escape(s[0])
                 for s in client.stream_list()
                 if s[0].startswith(prefix) )

    def layout(self, prefix, parsed_args, **kwargs):
        types = [ "int8", "int16", "int32", "int64",
                  "uint8", "uint16", "uint32", "uint64",
                  "float32", "float64" ]
        layouts = []
        for i in range(1,10):
            layouts.extend([(t + "_" + str(i)) for t in types])
        return ( l for l in layouts if l.startswith(prefix) )

    def meta_key(self, prefix, parsed_args, **kwargs):
        return (kv.split('=')[0] for kv
                in self.meta_keyval(prefix, parsed_args, **kwargs))

    def meta_keyval(self, prefix, parsed_args, **kwargs):
        client = nilmdb.client.Client(parsed_args.url)
        path = parsed_args.path
        if not path:
            return []
        results = []
        # prefix comes in as UTF-8, but results need to be Unicode,
        # weird.  Still doesn't work in all cases, but that's bugs in
        # argcomplete.
        prefix = nilmdb.utils.str.decode(prefix)
        for (k,v) in client.stream_get_metadata(path).items():
            kv = self.escape(k + '=' + v)
            if kv.startswith(prefix):
                results.append(kv)
        return results

class Cmdline(object):

    def __init__(self, argv = None):
        self.argv = argv or sys.argv[1:]
        try:
            # Assume command line arguments are encoded with stdin's encoding,
            # and reverse it.  Won't be needed in Python 3, but for now..
            self.argv = [ x.decode(sys.stdin.encoding) for x in self.argv ]
        except Exception: # pragma: no cover
            pass
        self.client = None
        self.def_url = os.environ.get("NILMDB_URL", "http://localhost/nilmdb/")
        self.subcmd = {}
        self.complete = Complete()

    def arg_time(self, toparse):
        """Parse a time string argument"""
        try:
            return nilmdb.utils.time.parse_time(toparse)
        except ValueError as e:
            raise argparse.ArgumentTypeError(sprintf("%s \"%s\"",
                                                     str(e), toparse))

    # Set up the parser
    def parser_setup(self):
        self.parser = JimArgumentParser(add_help = False,
                                        formatter_class = def_form)

        group = self.parser.add_argument_group("General options")
        group.add_argument("-h", "--help", action='help',
                           help='show this help message and exit')
        group.add_argument("-v", "--version", action="version",
                           version = nilmdb.__version__)

        group = self.parser.add_argument_group("Server")
        group.add_argument("-u", "--url", action="store",
                           default=self.def_url,
                           help="NilmDB server URL (default: %(default)s)"
                           ).completer = self.complete.url

        sub = self.parser.add_subparsers(
            title="Commands", dest="command",
            description="Use 'help command' or 'command --help' for more "
            "details on a particular command.")

        # Set up subcommands (defined in separate files)
        for cmd in subcommands:
            self.subcmd[cmd] = subcmd_mods[cmd].setup(self, sub)

    def die(self, formatstr, *args):
        fprintf(sys.stderr, formatstr + "\n", *args)
        if self.client:
            self.client.close()
        sys.exit(-1)

    def run(self):
        # Set SIGPIPE to its default handler -- we don't need Python
        # to catch it for us.
        try:
            signal.signal(signal.SIGPIPE, signal.SIG_DFL)
        except ValueError: # pragma: no cover
            pass

        # Clear cached timezone, so that we can pick up timezone changes
        # while running this from the test suite.
        datetime_tz._localtz = None

        # Run parser
        self.parser_setup()
        if argcomplete: # pragma: no cover
            argcomplete.autocomplete(self.parser)
        self.args = self.parser.parse_args(self.argv)

        # Run arg verify handler if there is one
        if "verify" in self.args:
            self.args.verify(self)

        self.client = nilmdb.client.Client(self.args.url)

        # Make a test connection to make sure things work,
        # unless the particular command requests that we don't.
        if "no_test_connect" not in self.args:
            try:
                server_version = self.client.version()
            except nilmdb.client.Error as e:
                self.die("error connecting to server: %s", str(e))

        # Now dispatch client request to appropriate function.  Parser
        # should have ensured that we don't have any unknown commands
        # here.
        retval = self.args.handler(self) or 0

        self.client.close()
        sys.exit(retval)