class AtomCLI: def __init__(self): self.element = Element(f"atom-cli_{uname().nodename}_{uuid4().hex}") self.indent = 2 self.style = Style.from_dict({ "logo_color": "#6039C8", }) self.session = PromptSession(style=self.style) self.print_atom_os_logo() self.serialization = "msgpack" self.cmd_map = { "help": self.cmd_help, "list": self.cmd_list, "records": self.cmd_records, "command": self.cmd_command, "read": self.cmd_read, "exit": self.cmd_exit, "serialization": self.cmd_serialization, } self.usage = { "cmd_help": cleandoc(""" Displays available commands and shows usage for commands. Usage: help [<command>]"""), "cmd_list": cleandoc(""" Displays available elements, streams, or commands. Can filter streams and commands based on element. Usage: list elements list streams [<element>] list commands [<element>]"""), "cmd_records": cleandoc(""" Displays log records or command and response records. Can filter records from the last N seconds or from certain elements. Usage: records log [<last_N_seconds>] [<element>...] records cmdres [<last_N_seconds>] <element>..."""), "cmd_command": cleandoc(""" Sends a command to an element and displays the response. Usage: command <element> <element_command> [<data>]"""), "cmd_read": cleandoc(""" Displays the entries of an element's stream. Can provide a rate to print the entries for ease of reading. Usage: read <element> <stream> [<rate_hz>]"""), "cmd_exit": cleandoc(""" Exits the atom-cli tool. Can also use the shortcut CTRL+D. Usage: exit"""), "cmd_serialization": cleandoc(""" Sets serialization/deserialization setting to either use msgpack, Apache arrow, or no (de)serialization. Defaults to msgpack serialization. This setting is overriden by deserialization keys received in stream. Usage: serialization (msgpack | arrow | none)"""), } def run(self): """ The main loop of the CLI. Reads the user input, verifies the command exists and calls the command. """ while True: try: inp = self.session.prompt( "\n> ", auto_suggest=AutoSuggestFromHistory()).split(" ") if not inp: continue command, args = inp[0], inp[1:] if command not in self.cmd_map.keys(): print("Invalid command. Type 'help' for valid commands.") else: self.cmd_map[command](*args) # Handle CTRL+C so user can break loops without exiting except KeyboardInterrupt: pass # Exit on CTRL+D except EOFError: self.cmd_exit() except Exception as e: print(str(type(e)) + " " + str(e)) def print_atom_os_logo(self): f = Figlet(font="slant") logo = f.renderText("ATOM OS") print(HTML(f"<logo_color>{logo}</logo_color>"), style=self.style) def format_record(self, record): """ Takes a record out of Redis, decodes the keys and values (if possible) and returns a formatted json string sorted by keys. """ formatted_record = {} for k, v in record.items(): if type(k) is bytes: k = k.decode() if not self.serialization: try: v = v.decode() except: v = str(v) formatted_record[k] = v sorted_record = {k: v for k, v in sorted( formatted_record.items(), key=lambda x: x[0])} try: ret = json.dumps(sorted_record, indent=self.indent) except TypeError as te: ret = sorted_record finally: return ret def cmd_help(self, *args): usage = self.usage["cmd_help"] if len(args) > 1: print(usage) print("\nToo many arguments to 'help'.") return if args: # Prints the usage of the command if args[0] in self.cmd_map.keys(): print(self.usage[f"cmd_{args[0]}"]) else: print(f"Command {args[0]} does not exist.") else: print("Try 'help <command>' for usage on a command") print("Available commands:") for command in self.cmd_map.keys(): print(f" {command}") def cmd_list(self, *args): usage = self.usage["cmd_list"] mode_map = { "elements": self.element.get_all_elements, "streams": self.element.get_all_streams, "commands": self.element.get_all_commands } if not args: print(usage) print("\n'list' must have an argument.") return mode = args[0] if mode not in mode_map.keys(): print(usage) print("\nInvalid argument to 'list'.") return if len(args) > 1 and mode == "elements": print(usage) print(f"\nInvalid number of arguments for command 'list elements'.") return if len(args) > 2: print(usage) print("\n'list' takes at most 2 arguments.") return items = mode_map[mode](*args[1:]) if not items: print(f"No {mode} exist.") return for item in items: print(item) def cmd_records(self, *args): usage = self.usage["cmd_records"] if not args: print(usage) print("\n'records' must have an argument.") return mode = args[0] # Check for start time if len(args) > 1 and args[1].isdigit(): ms = int(args[1]) * 1000 start_time = str(int(self.element._get_redis_timestamp()) - ms) elements = set(args[2:]) # If no start time, go from the very beginning else: start_time = "0" elements = set(args[1:]) if mode == "log": records = self.mode_log(start_time, elements) elif mode == "cmdres": if not elements: print(usage) print( "\nMust provide elements from which to get command response streams from.") return records = self.mode_cmdres(start_time, elements) else: print(usage) print("\nInvalid argument to 'records'.") return if not records: print("No records.") return for record in records: print(self.format_record(record)) def mode_log(self, start_time, elements): """ Reads the logs from Atom's log stream. Args: start_time (str): The time from which to start reading logs. elements (list): The elements on which to filter the logs for. """ records = [] all_records = self.element.entry_read_since( None, "log", start_time, serialization=None) for record in all_records: if not elements or record["element"].decode() in elements: record = {key: (value if isinstance(value, str) else value.decode( )) for key, value in record.items()} # Decode strings only which are required to records.append(record) return records def mode_cmdres(self, start_time, elements): """ Reads the command and response records from the provided elements. Args: start_time (str): The time from which to start reading logs. elements (list): The elements to get the command and response records from. """ streams, records = [], [] for element in elements: streams.append(self.element._make_response_id(element)) streams.append(self.element._make_command_id(element)) for stream in streams: cur_records = self.element.entry_read_since( None, stream, start_time, serialization=None) for record in cur_records: for key, value in record.items(): try: if not isinstance(value, str): value = value.decode() except: try: value = ser.deserialize(value, method=self.serialization) except: pass finally: record[key] = value record["type"], record["element"] = stream.split(":") records.append(record) return sorted(records, key=lambda x: (x["id"], x["type"])) def cmd_command(self, *args): usage = self.usage["cmd_command"] if len(args) < 2: print(usage) print("\nToo few arguments.") return element_name = args[0] command_name = args[1] if len(args) >= 3: data = str(" ".join(args[2:])) if self.serialization: try: data = json.loads(data) except: print("Received improperly formatted data!") return else: data = "" resp = self.element.command_send(element_name, command_name, data, serialize=(self.serialization is not None), deserialize=(self.serialization is not None), serialization=self.serialization) # shouldn't be used if it's None print(self.format_record(resp)) def cmd_read(self, *args): usage = self.usage["cmd_read"] if len(args) < 2: print(usage) print("\nToo few arguments.") return if len(args) > 3: print(usage) print("\nToo many arguments.") return if len(args) == 3: try: rate = float(args[2]) if rate < 0: raise ValueError() except ValueError: print("rate must be an float greater than 0.") return else: rate = None element_name, stream_name = args[:2] last_timestamp = None while True: start_time = time.time() entries = self.element.entry_read_n(element_name, stream_name, 1, deserialize=(self.serialization is not None), serialization=self.serialization) # shouldn't be used if it's None if not entries: print(f"No data from {element_name} {stream_name}.") return entry = entries[0] timestamp = entry["id"] # Only print the entry if it is different from the previous one if timestamp != last_timestamp: last_timestamp = timestamp print(self.format_record(entry)) if rate: time.sleep(max(1 / rate - (time.time() - start_time), 0)) def cmd_serialization(self, *args): usage = self.usage["cmd_serialization"] if (len(args) != 1): print(usage) print(f"\nPass one argument: {ser.Serializations.print_values()}.") return # Otherwise try to get the new setting if ser.is_valid_serialization(args[0].lower()): self.serialization = args[0].lower() if args[0].lower() != "none" else None else: print(f"\nArgument must be one of {ser.Serializations.print_values()}.") print("Current serialization status is {}".format(self.serialization)) def cmd_exit(*args): print("Exiting.") sys.exit()
"e": "waveform", "s": "serialized" } res = element.command_send("record", "start", data, serialize=True) time.sleep(TIME_FOR_WAVEFORM + 2) # Strings we'll recognize for the plotting commands. This is pretty # rudimentary and can be improved with some better parsing/processing/NLP sinx_strings = ["show sin", "show sign", "show sine"] cosx_strings = [ "show cos", "show cosine", "show coast", "show coats", "show cosign" ] tanx_strings = ["show tan", "showtime"] print("listening..") last_id = element._get_redis_timestamp() while True: entries = element.entry_read_since("voice", "string", last_id=last_id, block=1000) if entries: last_id = entries[0]["id"] #voice_string = entries[0]["data"].decode().strip().lower().replace("-", "").split(" ") voice_string = entries[0]["data"].decode().lower() print("Got voice string {}".format(voice_string)) if any(x in voice_string for x in sinx_strings): print("Plotting sinx")