path: root/cli/lib/cli/climanager.py

#
# Copyright (C) 2015 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""BDK CLI management."""


import argparse
import sys
import textwrap

from cli import clicommand
from cli import importutils
import error


class Error(error.Error):
    """Raised when CommandGroup directories and files don't match
    expectations.
    """


class CommandGroup(object):
    """A CLI Command Group."""

    def __init__(self, name, path, parent=None):
        self.name = name
        self.parent = parent
        self.path = path
        self.group_class = self.GetGroupClassType()

        if parent is None:
            self.parser = argparse.ArgumentParser(
                prog=name, description=self.group_class.__doc__)
        else:
            self.parser = self._AddSubparser(parent.subparsers, name,
                                             self.group_class)
        self.subparsers = self.parser.add_subparsers()

        self.FindCommands()

    def _AddSubparser(self, subparsers, name, command_class):
        """Adds a new subparser.

        Handles some common functionality such as help and description
        text generation.

        We expect Command and CommandGroup classes to provide help text in
        their docstring. The first line should be a single sentence overview,
        and additional text may be provided for a more complete description.

        Args:
            subparsers: the subparsers object to add the new subparser to.
            name: subparser name.
            command_class: subparser Command or CommandGroup class.

        Returns:
            The new subparser object.
        """
        # We need to do a little bit of additional work with the text to account
        # for the fact that all docstring lines except the first will be
        # indented some unknown amount.
        # For cross-platform support, keep the line endings with splitlines() so
        # we can easily join them together again in the same way.
        description = command_class.__doc__.splitlines(True)
        short_help = description[0].strip()
        description = description[0] + textwrap.dedent(''.join(description[1:]))
        return subparsers.add_parser(
            name, description=description, help=short_help,
            formatter_class=argparse.RawDescriptionHelpFormatter)

    def AddCommandGroup(self, name, path):
        CommandGroup(name, path, parent=self)

    def AddCommand(self, command_type):
        com_parser = self._AddSubparser(
            self.subparsers, command_type.__name__.lower(), command_type)
        com_parser.set_defaults(command_type=command_type)

        # Add arguments to the parser. REMAINDER must come last.
        command_type.AddDefaultArgs(com_parser)
        command_type.Args(com_parser)
        self.group_class.GroupArgs(com_parser)
        if command_type.remainder_arg:
            com_parser.add_argument(command_type.remainder_arg[0],
                                    nargs=argparse.REMAINDER,
                                    help=command_type.remainder_arg[1])
        else:
            # Add a hidden and unused remainder arg. This enables better error
            # messages by making parse_known_args() more forgiving instead of
            # just exiting immediately on failure, and ensures commands with and
            # without explicit remainder args error out the same ways.
            com_parser.add_argument('--_remainder_args',
                                    nargs=argparse.REMAINDER,
                                    help=argparse.SUPPRESS)

        # Update the command with its parentage and parser.
        command_type.set_group(self)
        command_type.set_parser(com_parser)

    def FindTypes(self, module, find_type):
        class_types = []
        for item in module.__dict__.values():
            if issubclass(type(item), type):
                if issubclass(item, find_type):
                    class_types.append(item)
        return class_types

    def FindCommands(self):
        count = 0
        # Iterate modules looking for commands.
        for module in importutils.ModuleIter(self.path):
            class_types = self.FindTypes(module, clicommand.Command)
            for class_type in class_types:
                self.AddCommand(class_type)
                count += 1
        if count == 0:
            # TODO(b/25951591): Change to use metric reporting Exception.
            raise Error('No commands in command group {}'.format(self.name))

    def GetGroupClassType(self):
        package = importutils.LoadPackage(self.path)
        class_types = self.FindTypes(package, clicommand.Group)
        if len(class_types) != 1:
            # TODO(b/25951591): Change to use metric reporting Exceptions.
            raise Error(
                'Expected exactly 1 class type for group {}. Found {}'.format(
                    self.name, class_types))
        return class_types[0]


class Cli(object):
    """CLI management.

    Attributes:
        args: the parsed arg namespace (after Execute() has been run).
    """

    def __init__(self, name, root_path):
        """Initializes the Cli object.

        Args:
            name: top-level CLI name.
            root_path: path to the package containing top-level commands.
        """
        self.name = name
        self.root_path = root_path
        self.root_group = CommandGroup(name, root_path)
        self.args = None

    def AddCommandGroup(self, name, path):
        """Adds a command group to the CLI.

        Args:
            name: command name.
            path: path to the package containing the grouped subcommands.
        """
        self.root_group.AddCommandGroup(name, path)

    def _ParseArgs(self, args=None):
        """Parses the commandline.

        Checks some corner cases to make argparse behavior a little more
        sane. Most significantly, when parsing a command that has set the
        remainder_arg class attribute, we do the following:
            1. Start the remainder args at the first unknown arg whether it's
               positional or optional; by default argparse only starts the
               remainder at a positional arg.
            2. Remove a leading '--' from the remainder if it was used.

        The result is that any of these will now work as expected:
            bdk build platform -j40
            bdk adb -s 12345 shell
            bdk adb -- -s 12345 shell

        Args:
            args: list of string args to parse, or None to use sys.argv.

        Returns:
            The parsed argument namespace.
        """
        if args is None:
            args = sys.argv[1:]

        # parse_args() has some corner cases that have undesired behavior (see
        # http://b/27795688), so instead we use parse_known_args() then manually
        # add any leftovers into the remainder argument.
        (parsed, leftover) = self.root_group.parser.parse_known_args(args)

        # parse_known_args() will skip unknown optionals to grab positionals or
        # known optionals, which can cause some problems e.g.:
        #     bdk adb -s foo bar ==> remainder=[foo bar], leftover=[-s]
        # Given this parsing, it's impossible to generically determine where -s
        # should be placed when passed through.
        # To avoid this, we find the first unknown arg whether it's positional
        # or optional, and then treat everything from there on as passthrough
        # args.
        if leftover and leftover != args[-len(leftover):]:
            # Iterate backwards over args until we've seen everything in
            # leftover, and that's our actual remainder.
            arg_counts = {arg: leftover.count(arg) for arg in leftover}
            i = len(args)
            for i in range(len(args) - 1, -1, -1):
                arg = args[i]
                count = arg_counts.get(arg)
                if count == 1:
                    arg_counts.pop(arg)
                elif count > 1:
                    arg_counts[arg] = count - 1
                if not arg_counts:
                    break

            # Re-parse while explicitly blocking off the actual remainder.
            parsed = self.root_group.parser.parse_args(args[:i])
            leftover = args[i:]

        if parsed.command_type.remainder_arg:
            # Add any leftovers to the remainder arg and remove '--' if it was
            # used.
            remainder_arg_name = parsed.command_type.remainder_arg[0]
            leftover = getattr(parsed, remainder_arg_name) + leftover
            if leftover and leftover[0] == '--':
                leftover.pop(0)
            setattr(parsed, remainder_arg_name, leftover)
        elif leftover:
            # If got remainder args but didn't want them, we have a bogus input.
            # Print an error message and exit.
            parsed.command_type.parser.error(
                'unknown arguments: {}'.format(' '.join(leftover)))

        return parsed

    def Execute(self, args=None):
        """Runs the selected command.

        Args:
            args: list of string args to parse, or None to use sys.argv.

        Returns:
            The result of the selected command's Run() function.

        Raises:
            TypeError: command Run() function returned a non-integer value.
        """
        self.args = self._ParseArgs(args)
        command = self.args.command_type()
        result = command.RunWithMetrics(self.args)

        # Enforce the return type with a useful message.
        if not isinstance(result, int):
            raise TypeError('Non-integer return code from "{0} {1}". '
                            'Please report this bug!'.format(
                                self.args.command_type.group().name,
                                self.args.command_type.__name__.lower()))
        return result