class VoteLabelCommand(Command): def __init__(self): super().__init__(name='vlabel', brief='manage vote labels') self.parser = PypsiArgParser() self.subcmd = self.parser.add_subparsers(help='subcmd', dest='subcmd', required=True) self.subcmd.add_parser('list', help='list all vote labels') create_cmd = self.subcmd.add_parser('new', help='create new vote label') create_cmd.add_argument('name', action='store', help='label name') def complete(self, shell, args, prefix): if len(args) == 1 and args[0].startswith('-'): completions = command_completer(self.parser, shell, args, prefix) else: completions = command_completer(self.subcmd, shell, args, prefix) return completions def run(self, shell, args): try: args = self.parser.parse_args(args) except CommandShortCircuit as err: return err.code if args.subcmd == 'list': return self.print_vote_labels(shell) if args.subcmd == 'new': return self.create_vote_label(shell, args.name) self.error(shell, f'unknown sub-command: {args.subcmd}') def print_vote_labels(self, shell): cursor = shell.ctx.db.vote_labels.find() table = Table( [Column('Id'), Column('Name'), Column('Symbol')], spacing=4) for row in cursor: table.append(row['_id'], row['name'], row['symbol']) table.write(sys.stdout) return 0 def create_vote_label(self, shell, name: str): vlabel = shell.ctx.db.create_vote_label(name) print('created new user successfully') print() print('Id: ', vlabel['_id']) print('Name: ', vlabel['name']) print('Symbol:', vlabel['symbol']) return 0
class ExchangeCommand(Command): def __init__(self): super().__init__(name='exchange', brief='manage exchanges') self.parser = PypsiArgParser() self.subcmd = self.parser.add_subparsers(help='subcmd', dest='subcmd') self.subcmd.add_parser('list', help='list all exchanges') create_cmd = self.subcmd.add_parser('new', help='create new exchange') create_cmd.add_argument('-s', '--select', action='store_true', help='select the exchange after creating it') create_cmd.add_argument('name', action='store', help='exchange name') select_cmd = self.subcmd.add_parser('select', help='select an active exhcnage') select_cmd.add_argument('id', help='exchange id') deselect_cmd = self.subcmd.add_parser('deselect', help='deselect active exchange') def complete(self, shell, args, prefix): if len(args) == 1 and args[0].startswith('-'): completions = command_completer(self.parser, shell, args, prefix) else: completions = command_completer(self.subcmd, shell, args, prefix) return completions def run(self, shell, args): try: args = self.parser.parse_args(args) except CommandShortCircuit as err: return err.code if args.subcmd == 'list': return self.print_exchanges(shell) if args.subcmd == 'new': return self.create_exchange(shell, args.name, args.select) if args.subcmd == 'select': return self.select_exchange(shell, args.id) if args.subcmd == 'deselect': shell.select_exchange(None) return 0 self.error(shell, f'unknown sub-command: {args.subcmd}') def print_exchanges(self, shell): table = Table([Column('Id'), Column('Name')], spacing=4) for row in shell.ctx.db.exchanges.find(): table.append(row['_id'], row['name']) table.write(sys.stdout) return 0 def create_exchange(self, shell, name: str, select: bool): exchange = shell.ctx.db.create_exchange(name) print('created new exchange successfully') print() print('Id: ', exchange['_id']) print('Name:', exchange['name']) if select: shell.select_exchange(exchange) return 0 def select_exchange(self, shell, exchange_id: str): exchange = shell.ctx.db.exchanges.find_one(ObjectId(exchange_id)) if not exchange: self.error(shell, f'exhcnage does not exist: {exchange_id}') return 1 shell.select_exchange(exchange) return 0
class StockCommand(Command): def __init__(self): super().__init__(name='stock', brief='manage stocks') self.parser = PypsiArgParser() self.subcmd = self.parser.add_subparsers(help='subcmd', dest='subcmd', required=True) list_cmd = self.subcmd.add_parser( 'list', help='list all stocks in an exchange') list_cmd.add_argument('-e', '--exchange-id', action='store', help='exchange id') create_cmd = self.subcmd.add_parser('new', help='create new stock') create_cmd.add_argument('name', action='store', help='stock name') create_cmd.add_argument('-s', '--symbol', action='store', help='stock symbol') create_cmd.add_argument('-e', '--exchange-id', action='store', help='exchange id') def complete(self, shell, args, prefix): if len(args) == 1 and args[0].startswith('-'): completions = command_completer(self.parser, shell, args, prefix) else: completions = command_completer(self.subcmd, shell, args, prefix) return completions def run(self, shell, args): try: args = self.parser.parse_args(args) except CommandShortCircuit as err: return err.code if args.subcmd == 'list': exchange_id = args.exchange_id if not exchange_id and shell.ctx.exchange: exchange_id = shell.ctx.exchange['_id'] return self.print_stocks(shell, exchange_id) if args.subcmd == 'new': exchange_id = args.exchange_id if not exchange_id and shell.ctx.exchange: exchange_id = shell.ctx.exchange['_id'] elif not exchange_id: self.error(shell, 'missing required argument -e/--exchange-id') return 1 return self.create_stock(shell, exchange_id, args.name, args.symbol) self.error(shell, f'unknown sub-command: {args.subcmd}') def print_stocks(self, shell, exchange_id: str = None): if exchange_id: columns = [ Column('Id'), Column('Name'), Column('Ups'), Column('Downs') ] cursor = shell.ctx.db.stocks.find( {'exchange_id': ObjectId(exchange_id)}) else: columns = [ Column('Id'), Column('Exchange Id'), Column('Name'), Column('Ups'), Column('Downs') ] cursor = shell.ctx.db.stocks.find() table = Table(columns, spacing=4) for row in cursor: if exchange_id: table.append(row['_id'], row['name'], row['ups'], row['downs']) else: table.append(row['_id'], row['exchange_id'], row['name'], row['ups'], row['downs']) table.write(sys.stdout) return 0 def create_stock(self, shell, exchange_id: str, name: str, symbol: str): stock = shell.ctx.db.create_stock(exchange_id, name, symbol) print('created new stock successfully') print() print('Id: ', stock['_id']) print('Name: ', stock['name']) print('Symbol:', stock['symbol']) return 0
class HistoryCommand(Command): ''' Interact with and manage the shell's history. ''' def __init__(self, name='history', brief='manage shell history', topic='shell', **kwargs): self.setup_parser(brief) super(HistoryCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def complete(self, shell, args, prefix): if len(args) == 1: return [ x for x in ('clear', 'delete', 'list', 'load', 'save') if x.startswith(prefix) ] if len(args) == 2: if args[0] == 'save' or args[0] == 'load': return path_completer(args[-1]) return [] def setup_parser(self, brief): self.parser = PypsiArgParser(prog='history', description=brief, usage=CmdUsage) subcmd = self.parser.add_subparsers(prog='history', dest='subcmd', metavar='subcmd') subcmd.required = True ls = subcmd.add_parser('list', help='list history events') ls.add_argument('count', metavar='N', type=int, help='number of events to display', nargs='?') subcmd.add_parser('clear', help='remove all history events') delete = subcmd.add_parser('delete', help='delete single history event') delete.add_argument( 'index', metavar='N', type=int, help='remove item at index N', ) save = subcmd.add_parser('save', help='save history to a file') save.add_argument('path', metavar='PATH', help='save history to file located at PATH') load = subcmd.add_parser('load', help='load history from a file') load.add_argument('path', metavar='PATH', help='load history from file located at PATH') def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code rc = 0 if ns.subcmd == 'list': start = 0 if ns.count and ns.count > 0: start = len(shell.ctx.history) - ns.count if start < 0: start = 0 i = start + 1 for event in shell.ctx.history[start:]: print(i, ' ', event, sep='') i += 1 elif ns.subcmd == 'clear': shell.ctx.history.clear() elif ns.subcmd == 'delete': try: del shell.ctx.history[ns.index - 1] except: self.error(shell, "invalid event index\n") rc = -1 elif ns.subcmd == 'save': try: with open(ns.path, 'w') as fp: for event in shell.ctx.history: fp.write(event) fp.write('\n') except IOError as e: self.error(shell, "error saving history to file: ", os.strerror(e.errno), '\n') rc = -1 elif ns.subcmd == 'load': try: lines = [] with safe_open(ns.path, 'r') as fp: for event in fp: lines.append(str(event)) shell.ctx.history.clear() for line in lines: shell.ctx.history.append(line.strip()) except IOError as e: self.error(shell, "error saving history to file: ", os.strerror(e.errno), '\n') rc = -1 except UnicodeEncodeError: self.error( shell, "error: file contains invalid unicode characters\n") return rc
class InputOutputCommand(Command): SUBCOMMAND_DELETE = "delete" SUBCOMMAND_EDIT = "edit" SUBCOMMAND_ADD = "add" def __init__(self, name, topic, brief, **kwargs): self.parser = PypsiArgParser(prog=name, description=brief) # TODO: start using args from ArgParser, now they only added to show pretty help self.parser = PypsiArgParser(prog=name, description=brief) subparsers = self.parser.add_subparsers(dest='subparser_name', help="Commands") add_arg_parser = subparsers.add_parser("add", help="add connection") edit_arg_parser = subparsers.add_parser("edit", help="edit connection") edit_arg_parser.add_argument("index", metavar="INDEX", action="store", help="connection index to edit") delete_arg_parser = subparsers.add_parser("delete", help="delete connection") delete_arg_parser.add_argument("index", metavar="INDEX", action="store", help="connection index to delete") super(InputOutputCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) self._handlers = { self.SUBCOMMAND_ADD: self._add_handler, self.SUBCOMMAND_DELETE: self._delete_handler, self.SUBCOMMAND_EDIT: self._edit_handler } self._print_header = True def run(self, shell, args): try: self.parser.parse_args(args) except CommandShortCircuit as e: return e.code component_info = shell._component_info self._print_header = shell._wizard_edit_mode if self.name == CommandName.INPUT: if component_info._inputs is None: component_info._inputs = [] self._inputs_or_outputs_array = component_info._inputs elif self.name == CommandName.OUTPUT: if component_info._outputs is None: component_info._outputs = [] self._inputs_or_outputs_array = component_info._outputs cmd = self.SUBCOMMAND_ADD if args is not None and len(args): if args[0] not in self._handlers.keys(): self.usage_error(shell, args) return 1 cmd = args[0] return self._handlers[cmd](shell, args) def _add_handler(self, shell, args=None): comp_conn = ComponentConnectionInfo() wizard_header = "{} Connection Configuration".format( self.name.capitalize()) ns = ConnectionWizard(wizard_header, comp_conn).run(shell, self._print_header) if ns: out_dict = Helper.filter_non_zero_string_items(ns.__dict__) comp_conn.load_from_json(out_dict) self._inputs_or_outputs_array.append(comp_conn) return 0 def _edit_handler(self, shell, args): if len(args) < 2: self.usage_error(shell, args) return 1 try: index = int(args[1]) except: print("Index parameter '{}' can not be casted to int".format( args[1])) self.usage_error(shell, args) return 0 if 0 <= index < len(self._inputs_or_outputs_array): comp_conn = self._inputs_or_outputs_array[index] wizard_header = "{} Connection Configuration".format( self.name.capitalize()) ns = ConnectionWizard(wizard_header, comp_conn).run(shell, self._print_header) if ns: out_dict = Helper.filter_non_zero_string_items(ns.__dict__) comp_conn.load_from_json(out_dict) else: print("Element with index '{}' is not found".format(index)) self.usage_error(shell, args) return 0 def _delete_handler(self, shell, args): if len(args) < 2: self.usage_error(shell, args) return 1 try: index = int(args[1]) except: print("Index parameter '{}' can not be casted to int".format( args[1])) self.usage_error(shell, args) return 0 if 0 <= index < len(self._inputs_or_outputs_array): self._inputs_or_outputs_array.pop(index) else: print("Element with index '{}' is not found".format(index)) self.usage_error(shell, args) return 0
class ArgumentCommand(Command): SUBCOMMAND_DELETE = "delete" SUBCOMMAND_EDIT = "edit" SUBCOMMAND_ADD = "add" def __init__(self, name, topic, brief, **kwargs): # TODO: start using args from ArgParser, now they only added to show pretty help self.parser = PypsiArgParser(prog=name, description=brief) subparsers = self.parser.add_subparsers(dest='subparser_name', help="Commands") add_arg_parser = subparsers.add_parser("add", help="add argument") edit_arg_parser = subparsers.add_parser("edit", help="edit argument") edit_arg_parser.add_argument("key", metavar="KEY", action="store", help="argument key to edit") delete_arg_parser = subparsers.add_parser("delete", help="delete argument") delete_arg_parser.add_argument("key", metavar="KEY", action="store", help="argument key to delete") super(ArgumentCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) self._handlers = { self.SUBCOMMAND_ADD: self._add_handler, self.SUBCOMMAND_DELETE: self._delete_handler, self.SUBCOMMAND_EDIT: self._edit_handler } def run(self, shell, args): try: self.parser.parse_args(args) except CommandShortCircuit as e: return e.code cmd = self.SUBCOMMAND_ADD self._print_header = shell._wizard_edit_mode if args is not None and len(args): if args[0] not in self._handlers.keys(): self.usage_error(shell, args) return 1 cmd = args[0] return self._handlers[cmd](shell, args) def _add_handler(self, shell, args=None): component_info = shell._component_info comp_arg = ComponentArgumentInfo() ns = ArgumentWizard(component_info, comp_arg).run(shell, self._print_header) if ns: out_dict = Helper.filter_non_zero_string_items(ns.__dict__) comp_arg.load_from_json(out_dict) if component_info.arguments is None: component_info.arguments = [] component_info.arguments.append(comp_arg) return 0 def _edit_handler(self, shell, args): if len(args) < 2: self.usage_error(shell, args) return 1 component_info = shell._component_info key = args[1] comp_arg = component_info.get_argument(key) if comp_arg: ns = ArgumentWizard(component_info, comp_arg).run(shell, self._print_header) if ns: out_dict = Helper.filter_non_zero_string_items(ns.__dict__) comp_arg.load_from_json(out_dict) else: print("Argument with the key '{}' is not found".format(key)) self.usage_error(shell, args) return 0 def _delete_handler(self, shell, args): if len(args) < 2: self.usage_error(shell, args) return 1 component_info = shell._component_info key = args[1] arg = component_info.get_argument(key) if arg: component_info.arguments.remove(arg) else: print("Argument with the key '{}' is not found".format(key)) self.usage_error(shell, args) return 0
class HistoryCommand(Command): """ Interact with and manage the shell's history. """ def __init__(self, name="history", brief="manage shell history", topic="shell", **kwargs): self.setup_parser(brief) super(HistoryCommand, self).__init__( name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs ) def complete(self, shell, args, prefix): if len(args) == 1: return [x for x in ("clear", "delete", "list", "load", "save") if x.startswith(prefix)] if len(args) == 2: if args[0] == "save" or args[0] == "load": return path_completer(args[-1]) return [] def setup_parser(self, brief): self.parser = PypsiArgParser(prog="history", description=brief, usage=CmdUsage) subcmd = self.parser.add_subparsers(prog="history", dest="subcmd", metavar="subcmd") subcmd.required = True ls = subcmd.add_parser("list", help="list history events") ls.add_argument("count", metavar="N", type=int, help="number of events to display", nargs="?") subcmd.add_parser("clear", help="remove all history events") delete = subcmd.add_parser("delete", help="delete single history event") delete.add_argument("index", metavar="N", type=int, help="remove item at index N") save = subcmd.add_parser("save", help="save history to a file") save.add_argument("path", metavar="PATH", help="save history to file located at PATH") load = subcmd.add_parser("load", help="load history from a file") load.add_argument("path", metavar="PATH", help="load history from file located at PATH") def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code rc = 0 if ns.subcmd == "list": start = 0 if ns.count: start = len(shell.ctx.history) - ns.count if start < 0: start = 0 i = start + 1 for event in shell.ctx.history[start:]: print(i, " ", event, sep="") i += 1 elif ns.subcmd == "clear": shell.ctx.history.clear() elif ns.subcmd == "delete": try: del shell.ctx.history[ns.index - 1] except: self.error(shell, "invalid event index\n") rc = -1 elif ns.subcmd == "save": try: with open(ns.path, "w") as fp: for event in shell.ctx.history: fp.write(event) fp.write("\n") except IOError as e: self.error(shell, "error saving history to file: ", os.strerror(e.errno), "\n") rc = -1 elif ns.subcmd == "load": try: lines = [] with safe_open(ns.path, "r") as fp: for event in fp: lines.append(str(event)) shell.ctx.history.clear() for line in lines: shell.ctx.history.append(line.strip()) except IOError as e: self.error(shell, "error saving history to file: ", os.strerror(e.errno), "\n") rc = -1 except UnicodeEncodeError: self.error(shell, "error: file contains invalid unicode characters\n") return rc
class HistoryCommand(Command): ''' Interact with and manage the shell's history. ''' def __init__(self, name='history', brief='manage shell history', topic='shell', **kwargs): self.setup_parser(brief) super(HistoryCommand, self).__init__( name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs ) def complete(self, shell, args, prefix): if len(args) == 1: return [x for x in ('clear', 'delete', 'list', 'load', 'save') if x.startswith(prefix)] if len(args) == 2: if args[0] == 'save' or args[0] == 'load': return path_completer(args[-1], prefix=prefix) return [] def setup_parser(self, brief): self.parser = PypsiArgParser( prog='history', description=brief, usage=CmdUsage ) subcmd = self.parser.add_subparsers(prog='history', dest='subcmd', metavar='subcmd') subcmd.required = True ls = subcmd.add_parser('list', help='list history events') ls.add_argument( 'count', metavar='N', type=int, help='number of events to display', nargs='?' ) subcmd.add_parser('clear', help='remove all history events') delete = subcmd.add_parser( 'delete', help='delete single history event' ) delete.add_argument( 'index', metavar='N', type=int, help='remove item at index N', ) save = subcmd.add_parser('save', help='save history to a file') save.add_argument( 'path', metavar='PATH', help='save history to file located at PATH' ) load = subcmd.add_parser('load', help='load history from a file') load.add_argument( 'path', metavar='PATH', help='load history from file located at PATH' ) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code rc = 0 if ns.subcmd == 'list': start = 0 if ns.count and ns.count > 0: start = len(shell.ctx.history) - ns.count if start < 0: start = 0 i = start + 1 for event in shell.ctx.history[start:]: print(i, ' ', event, sep='') i += 1 elif ns.subcmd == 'clear': shell.ctx.history.clear() elif ns.subcmd == 'delete': try: del shell.ctx.history[ns.index - 1] except: self.error(shell, "invalid event index\n") rc = -1 elif ns.subcmd == 'save': try: with open(ns.path, 'w') as fp: for event in shell.ctx.history: fp.write(event) fp.write('\n') except IOError as e: self.error(shell, "error saving history to file: ", os.strerror(e.errno), '\n') rc = -1 elif ns.subcmd == 'load': try: lines = [] with safe_open(ns.path, 'r') as fp: for event in fp: lines.append(str(event)) shell.ctx.history.clear() for line in lines: shell.ctx.history.append(line.strip()) except IOError as e: self.error(shell, "error saving history to file: ", os.strerror(e.errno), '\n') rc = -1 except UnicodeEncodeError: self.error( shell, "error: file contains invalid unicode characters\n" ) return rc
class VoteCommand(Command): def __init__(self): super().__init__(name='vote', brief='manage votes') self.parser = PypsiArgParser() subcmd = self.parser.add_subparsers(help='subcmd', dest='subcmd', required=True) list_cmd = subcmd.add_parser('list', help='list all stocks in an exchange') list_cmd.add_argument('-e', '--exchange-id', action='store', help='exchange id') create_cmd = subcmd.add_parser('new', help='create new stock') create_cmd.add_argument('name', action='store', help='stock name') create_cmd.add_argument('-e', '--exchange-id', action='store', help='exchange id', required=True) create_cmd.add_argument('-v', '--value', action='store', type=int, default=1000, help='initial stock value') def run(self, shell, args): try: args = self.parser.parse_args(args) except CommandShortCircuit as err: return err.code if args.subcmd == 'list': return self.print_stocks(shell, args.exchange_id) if args.subcmd == 'new': return self.create_stock(shell, args.exchange_id, args.name, args.value) self.error(shell, f'unknown sub-command: {args.subcmd}') def print_stocks(self, shell, exchange_id: str = None): if exchange_id: columns = [Column('Id'), Column('Name'), Column('Value')] cursor = shell.ctx.db.stocks.find( {'exchange_id': ObjectId(exchange_id)}) else: columns = [ Column('Id'), Column('Exchange Id'), Column('Name'), Column('Value') ] cursor = shell.ctx.db.stocks.find() table = Table(columns, spacing=4) for row in cursor: if exchange_id: table.append(row['_id'], row['name'], row['value']) else: table.append(row['_id'], row['exchange_id'], row['name'], row['value']) table.write(sys.stdout) return 0 def create_stock(self, shell, exchange_id: str, name: str, value: int): stock = shell.ctx.db.create_stock(exchange_id, name, value) print('created new stock successfully') print() print('Id: ', stock['_id']) print('Name: ', stock['name']) print('Value:', stock['value']) return 0
class MinionCommand(Command): def __init__(self, topic='proxy', name='minion', brief=HelpTopic, **kwargs): self.parser = PypsiArgParser(prog=name, description=brief, usage=MinionCmdUsage) subcmd = self.parser.add_subparsers(prog='minion', dest='subcmd', metavar='subcmd') subcmd.required = True shutdown = subcmd.add_parser('shutdown', help='shutdown the connected minion') restart = subcmd.add_parser('restart', help='restart the connected minion') load = subcmd.add_parser('load', help='loads a plugin into minion') load.add_argument('plugin', metavar='PLUGIN', help='name of plugin to load into minion') log = subcmd.add_parser('log', help='print the log of a minion') log.add_argument('lines', metavar='LINES', help='range of line number to print <start:finish>') super(MinionCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): if not self.check_connected(shell): self.error("shell", "Not currently connected to minion") return 0 try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code if ns.subcmd == 'shutdown': self.shutdown_minion(shell) elif ns.subcmd == 'restart': self.restart_minion(shell) elif ns.subcmd == 'load': self.load_plugin(shell, ns.plugin) def complete(self, shell, args, prefix): cmds = ['shutdown', 'restart', 'load', 'log'] if len(args) == 1: return sorted([cmd for cmd in cmds if cmd.startswith(prefix)]) def check_connected(self, shell): minion = shell.connect if minion is not None: return True else: False def shutdown_minion(self, shell): params = minionRPC_pb2.Empty() rpc = shell.minions[shell.connect]['rpc'] try: ret = rpc.shutdown(params) except Exception as e: shell.log.error(str(e)) self.error('shell', f"Failed to shutdown minion: {shell.connect}") if ret.reply == "okay": print(f'Shutting down minion: {shell.connect}') else: self.error('shell', f"Failed to shutdown minion: {shell.connect}") def restart_minion(self, shell): params = minionRPC_pb2.Empty() rpc = shell.minions[shell.connect]['rpc'] try: ret = rpc.runPluginFunction(params) except Exception as e: shell.log.error(str(e)) self.error('shell', f"Failed to restart minion: {shell.connect}") if ret.reply == "okay": print(f'Restarting minion: {shell.connect}') else: self.error('shell', f"Failed to restart minion: {shell.connect}") def load_plugin(self, shell, plugin): try: plg = LoaderUtil.load_plugin(plugin) except: self.error( 'shell', f"Plugin: {plugin} is not in the available plugin list") params = minionRPC_pb2.loadReq(plugin=plugin) rpc = shell.minions[shell.connect]['rpc'] try: ret = rpc.loadPlugin(params) except Exception as e: shell.log.error(str(e)) if ret.error < 1: print(f'Loaded plugin: {plugin} into minion: {shell.connect}') else: shell.log.error(ret.reply) self.error( 'shell', f"Failed to load plugin: { plugin } in {shell.connect}")
class UserCommand(Command): def __init__(self): super().__init__(name='user', brief='manage users') self.parser = PypsiArgParser() self.subcmd = self.parser.add_subparsers(help='subcmd', dest='subcmd', required=True) list_cmd = self.subcmd.add_parser('list', help='list all users in an exchange') list_cmd.add_argument('-e', '--exchange-id', action='store', help='exchange id') create_cmd = self.subcmd.add_parser('new', help='create new user') create_cmd.add_argument('name', action='store', help='user name') create_cmd.add_argument('-e', '--exchange-id', action='store', help='exchange id') describe_cmd = self.subcmd.add_parser('describe', help='describe user details') describe_cmd.add_argument('id', help='user id') def complete(self, shell, args, prefix): if len(args) == 1 and args[0].startswith('-'): completions = command_completer(self.parser, shell, args, prefix) else: completions = command_completer(self.subcmd, shell, args, prefix) return completions def run(self, shell, args): try: args = self.parser.parse_args(args) except CommandShortCircuit as err: return err.code if args.subcmd == 'list': exchange_id = args.exchange_id if not exchange_id and shell.ctx.exchange: exchange_id = shell.ctx.exchange['_id'] return self.print_users(shell, exchange_id) if args.subcmd == 'new': exchange_id = args.exchange_id if not exchange_id and shell.ctx.exchange: exchange_id = shell.ctx.exchange['_id'] elif not exchange_id: self.error(shell, 'missing required argument -e/--exchange-id') return 1 return self.create_user(shell, exchange_id, args.name) if args.subcmd == 'describe': return self.describe_user(shell, args.id) self.error(shell, f'unknown sub-command: {args.subcmd}') def print_users(self, shell, exchange_id: str = None): if exchange_id: columns = [Column('Id'), Column('Name')] cursor = shell.ctx.db.users.find( {'exchange_id': ObjectId(exchange_id)}) else: columns = [Column('Id'), Column('Exchange Id'), Column('Name')] cursor = shell.ctx.db.users.find() table = Table(columns, spacing=4) for row in cursor: if exchange_id: table.append(row['_id'], row['name']) else: table.append(row['_id'], row['exchange_id'], row['name']) table.write(sys.stdout) return 0 def create_user(self, shell, exchange_id: str, name: str): try: user = shell.ctx.db.create_user(exchange_id, name) except ItemDoesNotExist as err: self.error(shell, f'failed to create user: {err} - {exchange_id}') return -1 print('created new user successfully') print() print('Id: ', user['_id']) print('Name: ', user['name']) print('Token:', user['token']) return 0 def describe_user(self, shell, user_id: str): user = shell.ctx.db.users.find_one(ObjectId(user_id)) if not user: self.error(shell, 'user does not exist') return 1 print('Id: ', user['_id']) print('Exchange Id:', user['exchange_id']) print('Name: ', user['name']) print('Token: ', user['token']) return 0