Merge pull request #264 from lobocv/feature/submenu

tleonhardt · web-flow · commit 4895d5d8db4e · 2018-02-08T20:23:43.000-05:00
added a decorator function that allows nesting of Cmd classes to create submenus
diff --git a/cmd2.py b/cmd2.py
@@ -547,6 +547,230 @@ def strip_ansi(text):
     return ANSI_ESCAPE_RE.sub('', text)
 
 
+def _pop_readline_history(clear_history=True):
+    """Returns a copy of readline's history and optionally clears it (default)"""
+    history = [
+        readline.get_history_item(i)
+        for i in range(1, 1 + readline.get_current_history_length())
+    ]
+    if clear_history:
+        readline.clear_history()
+    return history
+
+
+def _push_readline_history(history, clear_history=True):
+    """Restores readline's history and optionally clears it first (default)"""
+    if clear_history:
+        readline.clear_history()
+    for line in history:
+        readline.add_history(line)
+
+
+def _complete_from_cmd(cmd_obj, text, line, begidx, endidx):
+    """Complete as though the user was typing inside cmd's cmdloop()"""
+    from itertools import takewhile
+    command_subcommand_params = line.split(None, 3)
+
+    if len(command_subcommand_params) < (3 if text else 2):
+        return cmd_obj.completenames(text)
+
+    command, subcommand = command_subcommand_params[:2]
+    n = len(command) + sum(1 for _ in takewhile(str.isspace, line))
+    cfun = getattr(cmd_obj, 'complete_' + subcommand, cmd_obj.complete)
+    return cfun(text, line[n:], begidx - n, endidx - n)
+
+
+class AddSubmenu(object):
+    """Conveniently add a submenu (Cmd-like class) to a Cmd
+
+    e.g. given "class SubMenu(Cmd): ..." then
+
+    @AddSubmenu(SubMenu(), 'sub')
+    class MyCmd(cmd.Cmd):
+        ....
+
+    will have the following effects:
+    1. 'sub' will interactively enter the cmdloop of a SubMenu instance
+    2. 'sub cmd args' will call do_cmd(args) in a SubMenu instance
+    3. 'sub ... [TAB]' will have the same behavior as [TAB] in a SubMenu cmdloop
+       i.e., autocompletion works the way you think it should
+    4. 'help sub [cmd]' will print SubMenu's help (calls its do_help())
+    """
+
+    class _Nonexistent(object):
+        """
+        Used to mark missing attributes.
+        Disable __dict__ creation since this class does nothing
+        """
+        __slots__ = ()  #
+
+    def __init__(self,
+                 submenu,
+                 command,
+                 aliases=(),
+                 reformat_prompt="{super_prompt}>> {sub_prompt}",
+                 shared_attributes=None,
+                 require_predefined_shares=True,
+                 create_subclass=False
+                 ):
+        """Set up the class decorator
+
+        submenu (Cmd):              Instance of something cmd.Cmd-like
+
+        command (str):              The command the user types to access the SubMenu instance
+
+        aliases (iterable):         More commands that will behave like "command"
+
+        reformat_prompt (str):      Format str or None to disable
+            if it's a string, it should contain one or more of:
+                {super_prompt}:     The current cmd's prompt
+                {command}:          The command in the current cmd with which it was called
+                {sub_prompt}:       The subordinate cmd's original prompt
+            the default is "{super_prompt}{command} {sub_prompt}"
+
+        shared_attributes (dict):   dict of the form {'subordinate_attr': 'parent_attr'}
+            the attributes are copied to the submenu at the last moment; the submenu's
+            attributes are backed up before this and restored afterward
+
+        require_predefined_shares: The shared attributes above must be independently
+            defined in the subordinate Cmd (default: True)
+
+        create_subclass: put the modifications in a subclass rather than modifying
+            the existing class (default: False)
+        """
+        self.submenu = submenu
+        self.command = command
+        self.aliases = aliases
+
+        if reformat_prompt is not None and not isinstance(reformat_prompt, str):
+            raise ValueError("reformat_prompt should be either a format string or None")
+        self.reformat_prompt = reformat_prompt
+
+        if require_predefined_shares:
+            for attr in shared_attributes.keys():
+                if not hasattr(submenu, attr):
+                    raise AttributeError("The shared attribute '{attr}' is not defined in {cmd}. Either define {attr} "
+                                         "in {cmd} or set require_predefined_shares=False."
+                                         .format(cmd=submenu.__class__.__name__, attr=attr))
+
+        self.shared_attributes = {} if shared_attributes is None else shared_attributes
+        self.create_subclass = create_subclass
+
+    def __call__(self, cmd_obj):
+        """Creates a subclass of Cmd wherein the given submenu can be accessed via the given command"""
+        def enter_submenu(parent_cmd, line):
+            """
+            This function will be bound to do_<submenu> and will change the scope of the CLI to that of the
+            submenu.
+            """
+            submenu = self.submenu
+            original_attributes = {attr: getattr(submenu, attr, AddSubmenu._Nonexistent)
+                                   for attr in self.shared_attributes.keys()
+                                   }
+            try:
+                # copy over any shared attributes
+                for sub_attr, par_attr in self.shared_attributes.items():
+                    setattr(submenu, sub_attr, getattr(parent_cmd, par_attr))
+
+                if line.parsed.args:
+                    # Remove the menu argument and execute the command in the submenu
+                    line = submenu.parser_manager.parsed(line.parsed.args)
+                    submenu.precmd(line)
+                    ret = submenu.onecmd(line)
+                    submenu.postcmd(ret, line)
+                else:
+                    history = _pop_readline_history()
+                    if self.reformat_prompt is not None:
+                        prompt = submenu.prompt
+                        submenu.prompt = self.reformat_prompt.format(
+                            super_prompt=parent_cmd.prompt,
+                            command=self.command,
+                            sub_prompt=prompt,
+                        )
+                    submenu.cmdloop()
+                    if self.reformat_prompt is not None:
+                        self.submenu.prompt = prompt
+                    _push_readline_history(history)
+            finally:
+                # copy back original attributes
+                for attr, value in original_attributes.items():
+                    if attr is not AddSubmenu._Nonexistent:
+                        setattr(submenu, attr, value)
+                    else:
+                        delattr(submenu, attr)
+
+        def complete_submenu(_self, text, line, begidx, endidx):
+            """
+            This function will be bound to complete_<submenu> and will perform the complete commands of the submenu.
+            """
+            submenu = self.submenu
+            original_attributes = {
+                attr: getattr(submenu, attr, AddSubmenu._Nonexistent)
+                for attr in self.shared_attributes.keys()
+            }
+            try:
+                # copy over any shared attributes
+                for sub_attr, par_attr in self.shared_attributes.items():
+                    setattr(submenu, sub_attr, getattr(_self, par_attr))
+
+                return _complete_from_cmd(submenu, text, line, begidx, endidx)
+            finally:
+                # copy back original attributes
+                for attr, value in original_attributes.items():
+                    if attr is not AddSubmenu._Nonexistent:
+                        setattr(submenu, attr, value)
+                    else:
+                        delattr(submenu, attr)
+
+        original_do_help = cmd_obj.do_help
+        original_complete_help = cmd_obj.complete_help
+
+        def help_submenu(_self, line):
+            """
+            This function will be bound to help_<submenu> and will call the help commands of the submenu.
+            """
+            tokens = line.split(None, 1)
+            if tokens and (tokens[0] == self.command or tokens[0] in self.aliases):
+                self.submenu.do_help(tokens[1] if len(tokens) == 2 else '')
+            else:
+                original_do_help(_self, line)
+
+        def _complete_submenu_help(_self, text, line, begidx, endidx):
+            """autocomplete to match help_submenu()'s behavior"""
+            tokens = line.split(None, 1)
+            if len(tokens) == 2 and (
+                    not (not tokens[1].startswith(self.command) and not any(
+                        tokens[1].startswith(alias) for alias in self.aliases))
+            ):
+                return self.submenu.complete_help(
+                    text,
+                    tokens[1],
+                    begidx - line.index(tokens[1]),
+                    endidx - line.index(tokens[1]),
+                )
+            else:
+                return original_complete_help(_self, text, line, begidx, endidx)
+
+        if self.create_subclass:
+            class _Cmd(cmd_obj):
+                do_help = help_submenu
+                complete_help = _complete_submenu_help
+        else:
+            _Cmd = cmd_obj
+            _Cmd.do_help = help_submenu
+            _Cmd.complete_help = _complete_submenu_help
+
+        # Create bindings in the parent command to the submenus commands.
+        setattr(_Cmd, 'do_'       + self.command, enter_submenu)
+        setattr(_Cmd, 'complete_' + self.command, complete_submenu)
+
+        # Create additional bindings for aliases
+        for _alias in self.aliases:
+            setattr(_Cmd, 'do_'       + _alias, enter_submenu)
+            setattr(_Cmd, 'complete_' + _alias, complete_submenu)
+        return _Cmd
+
+
 class Cmd(cmd.Cmd):
     """An easy but powerful framework for writing line-oriented command interpreters.
 
@@ -1386,6 +1610,7 @@ def do_shortcuts(self, _):
     def do_eof(self, _):
         """Called when <Ctrl>-D is pressed."""
         # End of script should not exit app, but <Ctrl>-D should.
+        print('')  # Required for clearing line when exiting submenu
         return self._STOP_AND_EXIT
 
     def do_quit(self, _):
diff --git a/examples/submenus.py b/examples/submenus.py
@@ -0,0 +1,113 @@
+#!/usr/bin/env python
+"""
+Create a CLI with a nested command structure as follows. The commands 'second' and 'third' navigate the CLI to the scope
+of the submenu. Nesting of the submenus is done with the cmd2.AddSubmenu() decorator.
+
+      (Top Level)----second----->(2nd Level)----third----->(3rd Level)
+        |                            |                          |
+        ---> say                     ---> say                   ---> say
+
+
+
+"""
+from __future__ import print_function
+import sys
+
+import cmd2
+from IPython import embed
+
+
+class ThirdLevel(cmd2.Cmd):
+    """To be used as a third level command class. """
+
+    def __init__(self, *args, **kwargs):
+        cmd2.Cmd.__init__(self, *args, **kwargs)
+        self.prompt = '3rdLevel '
+        self.top_level_attr = None
+        self.second_level_attr = None
+
+    def do_say(self, line):
+        print("You called a command in ThirdLevel with '%s'. " 
+              "It has access to top_level_attr: %s " 
+              "and second_level_attr: %s" % (line, self.top_level_attr, self.second_level_attr))
+
+    def help_say(self):
+        print("This is a third level submenu (submenu_ab). Options are qwe, asd, zxc")
+
+    def complete_say(self, text, line, begidx, endidx):
+        return [s for s in ['qwe', 'asd', 'zxc'] if s.startswith(text)]
+
+
+@cmd2.AddSubmenu(ThirdLevel(),
+                 command='third',
+                 aliases=('third_alias',),
+                 shared_attributes=dict(second_level_attr='second_level_attr', top_level_attr='top_level_attr'))
+class SecondLevel(cmd2.Cmd):
+    """To be used as a second level command class. """
+    def __init__(self, *args, **kwargs):
+        cmd2.Cmd.__init__(self, *args, **kwargs)
+        self.prompt = '2ndLevel '
+        self.top_level_attr = None
+        self.second_level_attr = 987654321
+
+    def do_ipy(self, arg):
+        """Enters an interactive IPython shell.
+
+        Run python code from external files with ``run filename.py``
+        End with ``Ctrl-D`` (Unix) / ``Ctrl-Z`` (Windows), ``quit()``, '`exit()``.
+        """
+        banner = 'Entering an embedded IPython shell type quit() or <Ctrl>-d to exit ...'
+        exit_msg = 'Leaving IPython, back to {}'.format(sys.argv[0])
+        embed(banner1=banner, exit_msg=exit_msg)
+
+    def do_say(self, line):
+        print("You called a command in SecondLevel with '%s'. "
+              "It has access to top_level_attr: %s" % (line, self.top_level_attr))
+
+    def help_say(self):
+        print("This is a SecondLevel menu. Options are qwe, asd, zxc")
+
+    def complete_say(self, text, line, begidx, endidx):
+        return [s for s in ['qwe', 'asd', 'zxc'] if s.startswith(text)]
+
+
+
+@cmd2.AddSubmenu(SecondLevel(),
+                 command='second',
+                 aliases=('second_alias',),
+                 shared_attributes=dict(top_level_attr='top_level_attr'))
+class TopLevel(cmd2.Cmd):
+    """To be used as the main / top level command class that will contain other submenus."""
+
+    def __init__(self, *args, **kwargs):
+        cmd2.Cmd.__init__(self, *args, **kwargs)
+        self.prompt = 'TopLevel '
+        self.top_level_attr = 123456789
+
+    def do_ipy(self, arg):
+        """Enters an interactive IPython shell.
+
+        Run python code from external files with ``run filename.py``
+        End with ``Ctrl-D`` (Unix) / ``Ctrl-Z`` (Windows), ``quit()``, '`exit()``.
+        """
+        banner = 'Entering an embedded IPython shell type quit() or <Ctrl>-d to exit ...'
+        exit_msg = 'Leaving IPython, back to {}'.format(sys.argv[0])
+        embed(banner1=banner, exit_msg=exit_msg)
+
+    def do_say(self, line):
+        print("You called a command in TopLevel with '%s'. "
+              "TopLevel has attribute top_level_attr=%s" % (line, self.top_level_attr))
+
+    def help_say(self):
+        print("This is a top level submenu. Options are qwe, asd, zxc")
+
+    def complete_say(self, text, line, begidx, endidx):
+        return [s for s in ['qwe', 'asd', 'zxc'] if s.startswith(text)]
+
+
+
+if __name__ == '__main__':
+
+    root = TopLevel()
+    root.cmdloop()
+
diff --git a/tests/test_submenu.py b/tests/test_submenu.py
