The CMD Module⌗

The cmd module is a really cool tool for Python development - it provides a lightweight wrapper for an interactive command-line interface. It’s quick and easy to add basic commands to your program, and you can give help messages using nothing more than function docstrings. Here’s an example:

from cmd import Cmd


class Shell(Cmd):
    prompt = "\nshell > "

    def do_exit(self, _) -> bool:
        """Exits the cmdloop."""
        return True

    def do_print(self, arg: str):
        """Prints your message."""
        print('Message:', arg)


def main():
    shell = Shell()
    shell.cmdloop()


if __name__ == "__main__":
    main()

And the output from the program:

$ python3.8 shell.py 

shell > help

Documented commands (type help <topic>):
========================================
exit  help  print


shell > help print
Prints your message.

shell > print a
Message: a

The Problem (and Solution)⌗

There are quite a few features that come with the CMD module, such as:

  • Help messages from docstrings, or optionally from help_* functions
  • Tab completion via completedefault and complete_* functions
  • Ability to hook pre and post-execution of commands

If it sounds interesting you should definitely read the documentation - it’s not very long and you’ll have a much better time making functional, good looking shells.

There is a major issue, however, that I seem to always encounter. Once you get past basic single-word commands, it becomes cumbersome and downright clunky to add support for command arguments and other options. The CMD module works with user input by definition, and anyone who has ever done input validation knows the feeling; it can take dozens of lines of string splitting, length checking, and type checking to even start actual functionality.

I know I’m not the first to experience this; a couple searches on PyPi for shell or cmd show some existing packages attempting to solve the problem. One in particular deserves a mention: cmd2. This is a project worth checking out, but after reading the documentation I realized the scope of this project is more broad than I intend to cover and involves quite a bit more code than I belive is necessary for my own purposes.

Fortunately, Python already provides a library specifically for this purpose: argparse. It’s not quite as easy as plug and play - the ArgumentParser expects tokenized arguments from sys.argv by default, will call sys.exit any time it runs into an error, and it requires careful planning to build your parsers to cooperate with cmd features like help messages.

The example provided below will walk through one solution for solving those issues and more, and my goal is to document support for integration of more features like tab completion and formatted help messages from argparse in future blog posts.

Example Solution⌗

The code below may be slightly intimidating at first, but I designed it so that there are only two additional requirements for each command in your shell: a function to build the required ArgumentParser and a decorator on the corresponding function.

import functools
import argparse
import shlex
import cmd
import sys


class ParserExitException(Exception):
    """Invoked when the parser calls sys.exit"""
    pass


class Parser(argparse.ArgumentParser):

    def exit(self, status=0, message=None):
        if message:
            self._print_message(message, sys.stderr)
        
        if status != 0:
            raise ParserExitExcepction


def parse_args(parse_func):
    def parse_args_decorator(func):
        parser = parse_func()
        parser.prog = func.__name__[3:]
        parser.description = func.__doc__

        @functools.wraps(func)
        def doer(shell, args):
            try:
                shell.namespace = parser.parse_args(shlex.split(args))
                result = func(shell, args)
                shell.namespace = None
                return result
            except ParserExitException:
                pass
        return doer
    return parse_args_decorator


def parse_print() -> Parser:
    parser = Parser(add_help=False)
    parser.add_argument('message')
    return parser


class Shell(cmd.Cmd):

    def __init__(self) -> None:
        super().__init__()
        self.prompt = '\nshell > '
        self.namespace = None

    @parse_args(parse_print)
    def do_print(self, _) -> None:
        """Prints the given message."""
        if self.namespace is not None:
            print(self.namespace.message)

    def do_exit(self, _):
        return True


if __name__ == '__main__':
    shell = Shell()
    shell.cmdloop()

Parser - ArgumentParser Subclass⌗

The first part of the program is a custom class inheriting from ArgumentParser.

class ParserExitExcepction(Exception):
    """Invoked when the parser calls sys.exit"""
    pass


class Parser(argparse.ArgumentParser):

    def exit(self, status=0, message=None):
        if message:
            self._print_message(message, sys.stderr)
        
        if status != 0:
            raise ParserExitExcepction

The only functionality that is overwritten from the parent class is the exit function, which is called whenever ArgumentParser encounters an error, when the help functionality is invoked, and in a few other instances. The safest way I saw to override this functionality was to raise a new exception instead of the default functionality: calling sys.exit(status).

Note: if you read the documentation you may notice a new parameter for ArgumentParser - exit_on_error. It may seem like this would solve the issue described above, but in practice it only prevents exiting during an ArgumentError and remains unchanged for calls to help and other situations.

Parser Function⌗

The next part I will discuss is the parse_* function.

def parse_print() -> Parser:
    parser = Parser(add_help=False)
    parser.add_argument('message')
    return parser

This function exists simply to instantiate and configure the Parser for a given do_* function. I tried several different methods of placement for this functionality, and ultimately the parser should be initialized outside the context of the do_* function if you want to use the built in HelpFormatter in your shell.

In order to associate the parser with the Shell class at instantiation, not at runtime, I pass this class as an argument to the decorator below.

Decorator⌗

The decorator was easily the most complex part of this program to design, but doesn’t actually do too much real work.

def parse_args(parse_func):
    def parse_args_decorator(func):
        parser = parse_func()
        parser.prog = func.__name__[3:]
        parser.description = func.__doc__

        @functools.wraps(func)
        def doer(shell, args):
            try:
                shell.namespace = parser.parse_args(shlex.split(args))
                result = func(shell, args)
                shell.namespace = None
                return result
            except ParserExitException:
                pass
        return doer
    return parse_args_decorator

As mentioned, this allows the Parser class to associate parsers with do_* functions at instantiation. This, coupled with the functionality provided by functools.wraps, allows me to set the parser’s program name to the command name and description to the command docstring. Now, argparse and cmd help messages are both linked to one another.

When wrapping the actual command function, I am able to catch ParserExitException to prevent program exiting, and as a hidden bonus I use shlex.split(args) to enable shell-like parsing of the command arguments to include support for things like quoting (e.g. print "a b c").

The Shell⌗

Finally, the actual Shell class.

class Shell(cmd.Cmd):

    def __init__(self) -> None:
        super().__init__()
        self.prompt = '\nshell > '
        self.namespace = None

    @parse_args(parse_print)
    def do_print(self, _) -> None:
        """Prints the given message."""
        if self.namespace is not None:
            print(self.namespace.message)

    def do_exit(self, _):
        return True

As you can see, not much has changed. There are only two key differences: commands with parsable arguments require the @parse_args decorator, and in order to access those parsed arguments functions should use the self.namespace variable.

My objective was to make this process simple to expand, and I belive that much has been accomplished. At the very least, I hope you now know a little bit more about both cmd and argparse. In the next post on this topic, I plan to build on this example and add support for tab completion as well as full integration of help messages.