#!/usr/bin/env python # coding=utf-8 """A simple example demonstrating how to use flag and index based tab-completion functions """ import argparse import AutoCompleter import itertools from typing import List import cmd2 from cmd2 import with_argparser, with_category class TabCompleteExample(cmd2.Cmd): """ Example cmd2 application where we a base command which has a couple subcommands.""" CAT_AUTOCOMPLETE = 'AutoComplete Examples' def __init__(self): super().__init__() # For mocking a data source for the example commands ratings_types = ['G', 'PG', 'PG-13', 'R', 'NC-17'] show_ratings = ['TV-Y', 'TV-Y7', 'TV-G', 'TV-PG', 'TV-14', 'TV-MA'] static_list_directors = ['J. J. Abrams', 'Irvin Kershner', 'George Lucas', 'Richard Marquand', 'Rian Johnson', 'Gareth Edwards'] actors = ['Mark Hamill', 'Harrison Ford', 'Carrie Fisher', 'Alec Guinness', 'Peter Mayhew', 'Anthony Daniels', 'Adam Driver', 'Daisy Ridley', 'John Boyega', 'Oscar Isaac', 'Lupita Nyong\'o', 'Andy Serkis', 'Liam Neeson', 'Ewan McGregor', 'Natalie Portman', 'Jake Lloyd', 'Hayden Christensen', 'Christopher Lee'] USER_MOVIE_LIBRARY = ['ROGUE1', 'SW_EP04', 'SW_EP05'] MOVIE_DATABASE_IDS = ['SW_EP01', 'SW_EP02', 'SW_EP03', 'ROGUE1', 'SW_EP04', 'SW_EP05', 'SW_EP06', 'SW_EP07', 'SW_EP08', 'SW_EP09'] MOVIE_DATABASE = {'SW_EP04': {'title': 'Star Wars: Episode IV - A New Hope', 'rating': 'PG', 'director': ['George Lucas'], 'actor': ['Mark Hamill', 'Harrison Ford', 'Carrie Fisher', 'Alec Guinness', 'Peter Mayhew', 'Anthony Daniels'] }, 'SW_EP05': {'title': 'Star Wars: Episode V - The Empire Strikes Back', 'rating': 'PG', 'director': ['Irvin Kershner'], 'actor': ['Mark Hamill', 'Harrison Ford', 'Carrie Fisher', 'Alec Guinness', 'Peter Mayhew', 'Anthony Daniels'] }, 'SW_EP06': {'title': 'Star Wars: Episode IV - A New Hope', 'rating': 'PG', 'director': ['Richard Marquand'], 'actor': ['Mark Hamill', 'Harrison Ford', 'Carrie Fisher', 'Alec Guinness', 'Peter Mayhew', 'Anthony Daniels'] }, 'SW_EP01': {'title': 'Star Wars: Episode I - The Phantom Menace', 'rating': 'PG', 'director': ['George Lucas'], 'actor': ['Liam Neeson', 'Ewan McGregor', 'Natalie Portman', 'Jake Lloyd'] }, 'SW_EP02': {'title': 'Star Wars: Episode II - Attack of the Clones', 'rating': 'PG', 'director': ['George Lucas'], 'actor': ['Liam Neeson', 'Ewan McGregor', 'Natalie Portman', 'Hayden Christensen', 'Christopher Lee'] }, 'SW_EP03': {'title': 'Star Wars: Episode III - Revenge of the Sith', 'rating': 'PG-13', 'director': ['George Lucas'], 'actor': ['Liam Neeson', 'Ewan McGregor', 'Natalie Portman', 'Hayden Christensen'] }, } USER_SHOW_LIBRARY = {'SW_REB': ['S01E01', 'S02E02']} SHOW_DATABASE_IDS = ['SW_CW', 'SW_TCW', 'SW_REB'] SHOW_DATABASE = {'SW_CW': {'title': 'Star Wars: Clone Wars', 'rating': 'TV-Y7', 'seasons': {1: ['S01E01', 'S01E02', 'S01E03'], 2: ['S02E01', 'S02E02', 'S02E03']} }, 'SW_TCW': {'title': 'Star Wars: The Clone Wars', 'rating': 'TV-PG', 'seasons': {1: ['S01E01', 'S01E02', 'S01E03'], 2: ['S02E01', 'S02E02', 'S02E03']} }, 'SW_REB': {'title': 'Star Wars: Rebels', 'rating': 'TV-Y7', 'seasons': {1: ['S01E01', 'S01E02', 'S01E03'], 2: ['S02E01', 'S02E02', 'S02E03']} }, } # This demonstrates a number of customizations of the AutoCompleter version of ArgumentParser # - The help output will separately group required vs optional flags # - The help output for arguments with multiple flags or with append=True is more concise # - ACArgumentParser adds the ability to specify ranges of argument counts in 'nargs' suggest_parser = AutoCompleter.ACArgumentParser() suggest_parser.add_argument('-t', '--type', choices=['movie', 'show'], required=True) suggest_parser.add_argument('-d', '--duration', nargs=(1, 2), action='append', help='Duration constraint in minutes.\n' '\tsingle value - maximum duration\n' '\t[a, b] - duration range') @with_category(CAT_AUTOCOMPLETE) @with_argparser(suggest_parser) def do_suggest(self, args) -> None: """Suggest command demonstrates argparse customizations See hybrid_suggest and orig_suggest to compare the help output. """ if not args.type: self.do_help('suggest') def complete_suggest(self, text: str, line: str, begidx: int, endidx: int) -> List[str]: """ Adds tab completion to media""" completer = AutoCompleter.AutoCompleter(TabCompleteExample.suggest_parser, 1) tokens, _ = self.tokens_for_completion(line, begidx, endidx) results = completer.complete_command(tokens, text, line, begidx, endidx) return results # If you prefer the original argparse help output but would like narg ranges, it's possible # to enable narg ranges without the help changes using this method suggest_parser_hybrid = argparse.ArgumentParser() # This registers the custom narg range handling AutoCompleter.register_custom_actions(suggest_parser_hybrid) suggest_parser_hybrid.add_argument('-t', '--type', choices=['movie', 'show'], required=True) suggest_parser_hybrid.add_argument('-d', '--duration', nargs=(1, 2), action='append', help='Duration constraint in minutes.\n' '\tsingle value - maximum duration\n' '\t[a, b] - duration range') @with_category(CAT_AUTOCOMPLETE) @with_argparser(suggest_parser_hybrid) def do_hybrid_suggest(self, args): if not args.type: self.do_help('orig_suggest') def complete_hybrid_suggest(self, text, line, begidx, endidx): """ Adds tab completion to media""" completer = AutoCompleter.AutoCompleter(TabCompleteExample.suggest_parser_hybrid) tokens, _ = self.tokens_for_completion(line, begidx, endidx) results = completer.complete_command(tokens, text, line, begidx, endidx) return results # This variant demonstrates the AutoCompleter working with the orginial argparse. # Base argparse is unable to specify narg ranges. Autocompleter will keep expecting additional arguments # for the -d/--duration flag until you specify a new flaw or end the list it with '--' suggest_parser_orig = argparse.ArgumentParser() suggest_parser_orig.add_argument('-t', '--type', choices=['movie', 'show'], required=True) suggest_parser_orig.add_argument('-d', '--duration', nargs='+', action='append', help='Duration constraint in minutes.\n' '\tsingle value - maximum duration\n' '\t[a, b] - duration range') @with_argparser(suggest_parser_orig) @with_category(CAT_AUTOCOMPLETE) def do_orig_suggest(self, args) -> None: if not args.type: self.do_help('orig_suggest') def complete_orig_suggest(self, text, line, begidx, endidx) -> List[str]: """ Adds tab completion to media""" completer = AutoCompleter.AutoCompleter(TabCompleteExample.suggest_parser_orig) tokens, _ = self.tokens_for_completion(line, begidx, endidx) results = completer.complete_command(tokens, text, line, begidx, endidx) return results ################################################################################### # The media command demonstrates a completer with multiple layers of subcommands # - This example uses a flat completion lookup dictionary def query_actors(self) -> List[str]: """Simulating a function that queries and returns a completion values""" return TabCompleteExample.actors def _do_media_movies(self, args) -> None: if not args.command: self.do_help('media movies') elif args.command == 'list': for movie_id in TabCompleteExample.MOVIE_DATABASE: movie = TabCompleteExample.MOVIE_DATABASE[movie_id] print('{}\n-----------------------------\n{} ID: {}\nDirector: {}\nCast:\n {}\n\n' .format(movie['title'], movie['rating'], movie_id, ', '.join(movie['director']), '\n '.join(movie['actor']))) def _do_media_shows(self, args) -> None: if not args.command: self.do_help('media shows') elif args.command == 'list': for show_id in TabCompleteExample.SHOW_DATABASE: show = TabCompleteExample.SHOW_DATABASE[show_id] print('{}\n-----------------------------\n{} ID: {}' .format(show['title'], show['rating'], show_id)) for season in show['seasons']: ep_list = show['seasons'][season] print(' Season {}:\n {}' .format(season, '\n '.join(ep_list))) print() media_parser = AutoCompleter.ACArgumentParser(prog='media') media_types_subparsers = media_parser.add_subparsers(title='Media Types', dest='type') movies_parser = media_types_subparsers.add_parser('movies') movies_parser.set_defaults(func=_do_media_movies) movies_commands_subparsers = movies_parser.add_subparsers(title='Commands', dest='command') movies_list_parser = movies_commands_subparsers.add_parser('list') movies_list_parser.add_argument('-t', '--title', help='Title Filter') movies_list_parser.add_argument('-r', '--rating', help='Rating Filter', nargs='+', choices=ratings_types) movies_list_parser.add_argument('-d', '--director', help='Director Filter') movies_list_parser.add_argument('-a', '--actor', help='Actor Filter', action='append') movies_add_parser = movies_commands_subparsers.add_parser('add') movies_add_parser.add_argument('title', help='Movie Title') movies_add_parser.add_argument('rating', help='Movie Rating', choices=ratings_types) movies_add_parser.add_argument('-d', '--director', help='Director', nargs=(1, 2), required=True) movies_add_parser.add_argument('actor', help='Actors', nargs='*') movies_delete_parser = movies_commands_subparsers.add_parser('delete') shows_parser = media_types_subparsers.add_parser('shows') shows_parser.set_defaults(func=_do_media_shows) shows_commands_subparsers = shows_parser.add_subparsers(title='Commands', dest='command') shows_list_parser = shows_commands_subparsers.add_parser('list') @with_category(CAT_AUTOCOMPLETE) @with_argparser(media_parser) def do_media(self, args): """Media management command demonstrates multiple layers of subcommands being handled by AutoCompleter""" func = getattr(args, 'func', None) if func is not None: # Call whatever subcommand function was selected func(self, args) else: # No subcommand was provided, so call help self.do_help('media') # This completer is implemented using a single dictionary to look up completion lists for all layers of # subcommands. For each argument, AutoCompleter will search for completion values from the provided # arg_choices dict. This requires careful naming of argparse arguments so that there are no unintentional # name collisions. def complete_media(self, text, line, begidx, endidx): """ Adds tab completion to media""" choices = {'actor': self.query_actors, # function 'director': TabCompleteExample.static_list_directors # static list } completer = AutoCompleter.AutoCompleter(TabCompleteExample.media_parser, arg_choices=choices) tokens, _ = self.tokens_for_completion(line, begidx, endidx) results = completer.complete_command(tokens, text, line, begidx, endidx) return results ################################################################################### # The library command demonstrates a completer with multiple layers of subcommands # with different completion results per sub-command # - This demonstrates how to build a tree of completion lookups to pass down # # Only use this method if you absolutely need to as it dramatically # increases the complexity and decreases readability. def _do_library_movie(self, args): if not args.type or not args.command: self.do_help('library movie') def _do_library_show(self, args): if not args.type: self.do_help('library show') def _query_movie_database(self): return list(set(TabCompleteExample.MOVIE_DATABASE_IDS).difference(set(TabCompleteExample.USER_MOVIE_LIBRARY))) def _query_movie_user_library(self): return TabCompleteExample.USER_MOVIE_LIBRARY def _filter_library(self, text, line, begidx, endidx, full, exclude=[]): candidates = list(set(full).difference(set(exclude))) return [entry for entry in candidates if entry.startswith(text)] library_parser = AutoCompleter.ACArgumentParser(prog='library') library_subcommands = library_parser.add_subparsers(title='Media Types', dest='type') library_movie_parser = library_subcommands.add_parser('movie') library_movie_parser.set_defaults(func=_do_library_movie) library_movie_subcommands = library_movie_parser.add_subparsers(title='Command', dest='command') library_movie_add_parser = library_movie_subcommands.add_parser('add') library_movie_add_parser.add_argument('movie_id', help='ID of movie to add', action='append') library_movie_add_parser.add_argument('-b', '--borrowed', action='store_true') library_movie_remove_parser = library_movie_subcommands.add_parser('remove') library_movie_remove_parser.add_argument('movie_id', help='ID of movie to remove', action='append') library_show_parser = library_subcommands.add_parser('show') library_show_parser.set_defaults(func=_do_library_show) library_show_subcommands = library_show_parser.add_subparsers(title='Command', dest='command') library_show_add_parser = library_show_subcommands.add_parser('add') library_show_add_parser.add_argument('show_id', help='Show IDs to add') library_show_add_parser.add_argument('episode_id', nargs='*', help='Show IDs to add') library_show_rmv_parser = library_show_subcommands.add_parser('remove') # Demonstrates a custom completion function that does more with the command line than is # allowed by the standard completion functions def _filter_episodes(self, text, line, begidx, endidx, show_db, user_lib): tokens, _ = self.tokens_for_completion(line, begidx, endidx) show_id = tokens[3] if show_id: if show_id in show_db: show = show_db[show_id] all_episodes = itertools.chain(*(show['seasons'].values())) if show_id in user_lib: user_eps = user_lib[show_id] else: user_eps = [] return self._filter_library(text, line, begidx, endidx, all_episodes, user_eps) return [] @with_category(CAT_AUTOCOMPLETE) @with_argparser(library_parser) def do_library(self, args): """Media management command demonstrates multiple layers of subcommands being handled by AutoCompleter""" func = getattr(args, 'func', None) if func is not None: # Call whatever subcommand function was selected func(self, args) else: # No subcommand was provided, so call help self.do_help('library') def complete_library(self, text, line, begidx, endidx): # this demonstrates the much more complicated scenario of having # unique completion parameters per sub-command that use the same # argument name. To do this we build a multi-layer nested tree # of lookups far AutoCompleter to traverse. This nested tree must # match the structure of the argparse parser # movie_add_choices = {'movie_id': self._query_movie_database} movie_remove_choices = {'movie_id': self._query_movie_user_library} # This demonstrates the ability to mix custom completion functions with argparse completion. # By specifying a tuple for a completer, AutoCompleter expects a custom completion function # with optional index-based as well as keyword based arguments. This is an alternative to using # a partial function. show_add_choices = {'show_id': (self._filter_library, # This is a custom completion function # This tuple represents index-based args to append to the function call (list(TabCompleteExample.SHOW_DATABASE.keys()),) ), 'episode_id': (self._filter_episodes, # this is a custom completion function # this list represents index-based args to append to the function call [TabCompleteExample.SHOW_DATABASE], # this dict contains keyword-based args to append to the function call {'user_lib': TabCompleteExample.USER_SHOW_LIBRARY})} show_remove_choices = {} # The library movie sub-parser group 'command' has 2 sub-parsers: # 'add' and 'remove' library_movie_command_params = \ {'add': (movie_add_choices, None), 'remove': (movie_remove_choices, None)} library_show_command_params = \ {'add': (show_add_choices, None), 'remove': (show_remove_choices, None)} # The 'library movie' command has a sub-parser group called 'command' library_movie_subcommand_groups = {'command': library_movie_command_params} library_show_subcommand_groups = {'command': library_show_command_params} # Mapping of a specific sub-parser of the 'type' group to a tuple. Each # tuple has 2 values corresponding what's passed to the constructor # parameters (arg_choices,subcmd_args_lookup) of the nested # instance of AutoCompleter library_type_params = {'movie': (None, library_movie_subcommand_groups), 'show': (None, library_show_subcommand_groups)} # maps the a subcommand group to a dictionary mapping a specific # sub-command to a tuple of (arg_choices, subcmd_args_lookup) # # In this example, 'library_parser' has a sub-parser group called 'type' # under the type sub-parser group, there are 2 sub-parsers: 'movie', 'show' library_subcommand_groups = {'type': library_type_params} completer = AutoCompleter.AutoCompleter(TabCompleteExample.library_parser, subcmd_args_lookup=library_subcommand_groups) tokens, _ = self.tokens_for_completion(line, begidx, endidx) results = completer.complete_command(tokens, text, line, begidx, endidx) return results if __name__ == '__main__': app = TabCompleteExample() app.cmdloop()