added a decorator function that allows nesting of Cmd instances to create submenus

1 files changed, 224 insertions, 0 deletions

diff --git a/cmd2.py b/cmd2.py
index 73ad491d..a3138c4c 100755
--- a/cmd2.py
+++ b/cmd2.py
@@ -438,6 +438,229 @@ 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, 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.completenames(text)
+
+    command, subcommand = command_subcommand_params[:2]
+    n = len(command) + sum(1 for _ in takewhile(str.isspace, line))
+    cfun = getattr(cmd, 'complete_' + subcommand, cmd.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={},
+                 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 = shared_attributes
+
+        self.create_subclass = create_subclass
+
+    def __call__(self, Cmd):
+        """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:
+                    # Execute the command
+                    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.do_help
+        original_complete_help = Cmd.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 (
+                tokens[1].startswith(self.command) or
+                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):
+                do_help = help_submenu
+                complete_help = _complete_submenu_help
+        else:
+            _Cmd = Cmd
+            _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.
 
@@ -1188,6 +1411,7 @@ class Cmd(cmd.Cmd):
     def do_eof(self, arg):
         """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, arg):