1# Copyright 2021 The Pigweed Authors 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); you may not 4# use this file except in compliance with the License. You may obtain a copy of 5# the License at 6# 7# https://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT 11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the 12# License for the specific language governing permissions and limitations under 13# the License. 14"""Utilities for creating an interactive console.""" 15 16from collections import defaultdict 17from itertools import chain 18import inspect 19import textwrap 20import types 21from typing import Any, Collection, Dict, Iterable, Mapping, NamedTuple 22 23import pw_status 24from pw_protobuf_compiler import python_protos 25 26import pw_rpc 27from pw_rpc.descriptors import Method 28from pw_rpc.console_tools import functions 29 30_INDENT = ' ' 31 32 33class CommandHelper: 34 """Used to implement a help command in an RPC console.""" 35 @classmethod 36 def from_methods(cls, 37 methods: Iterable[Method], 38 variables: Mapping[str, object], 39 header: str, 40 footer: str = '') -> 'CommandHelper': 41 return cls({m.full_name: m 42 for m in methods}, variables, header, footer) 43 44 def __init__(self, 45 methods: Mapping[str, object], 46 variables: Mapping[str, object], 47 header: str, 48 footer: str = ''): 49 self._methods = methods 50 self._variables = variables 51 self.header = header 52 self.footer = footer 53 54 def help(self, item: object = None) -> str: 55 """Returns a help string with a command or all commands listed.""" 56 57 if item is None: 58 all_vars = '\n'.join(sorted(self._variables_without_methods())) 59 all_rpcs = '\n'.join(self._methods) 60 return (f'{self.header}\n\n' 61 f'All variables:\n\n{textwrap.indent(all_vars, _INDENT)}' 62 '\n\n' 63 f'All commands:\n\n{textwrap.indent(all_rpcs, _INDENT)}' 64 f'\n\n{self.footer}'.strip()) 65 66 # If item is a string, find commands matching that. 67 if isinstance(item, str): 68 matches = {n: m for n, m in self._methods.items() if item in n} 69 if not matches: 70 return f'No matches found for {item!r}' 71 72 if len(matches) == 1: 73 name, method = next(iter(matches.items())) 74 return f'{name}\n\n{inspect.getdoc(method)}' 75 76 return f'Multiple matches for {item!r}:\n\n' + textwrap.indent( 77 '\n'.join(matches), _INDENT) 78 79 return inspect.getdoc(item) or f'No documentation for {item!r}.' 80 81 def _variables_without_methods(self) -> Mapping[str, object]: 82 packages = frozenset( 83 n.split('.', 1)[0] for n in self._methods if '.' in n) 84 85 return { 86 name: var 87 for name, var in self._variables.items() if name not in packages 88 } 89 90 def __call__(self, item: object = None) -> None: 91 """Prints the help string.""" 92 print(self.help(item)) 93 94 def __repr__(self) -> str: 95 """Returns the help, so foo and foo() are equivalent in a console.""" 96 return self.help() 97 98 99class ClientInfo(NamedTuple): 100 """Information about an RPC client as it appears in the console.""" 101 # The name to use in the console to refer to this client. 102 name: str 103 104 # An object to use in the console for the client. May be a pw_rpc.Client. 105 client: object 106 107 # The pw_rpc.Client; may be the same object as client. 108 rpc_client: pw_rpc.Client 109 110 111class Context: 112 """The Context class is used to set up an interactive RPC console. 113 114 The Context manages a set of variables that make it easy to access RPCs and 115 protobufs in a REPL. 116 117 As an example, this class can be used to set up a console with IPython: 118 119 .. code-block:: python 120 121 context = console_tools.Context( 122 clients, default_client, protos, help_header=WELCOME_MESSAGE) 123 IPython.terminal.embed.InteractiveShellEmbed().mainloop( 124 module=types.SimpleNamespace(**context.variables())) 125 """ 126 def __init__(self, 127 client_info: Collection[ClientInfo], 128 default_client: Any, 129 protos: python_protos.Library, 130 *, 131 help_header: str = '') -> None: 132 """Creates an RPC console context. 133 134 Protos and RPC services are accessible by their proto package and name. 135 The target for these can be set with the set_target function. 136 137 Args: 138 client_info: ClientInfo objects that represent the clients this 139 console uses to communicate with other devices 140 default_client: default client object; must be one of the clients 141 protos: protobufs to use for RPCs for all clients 142 help_header: Message to display for the help command 143 """ 144 assert client_info, 'At least one client must be provided!' 145 146 self.client_info = client_info 147 self.current_client = default_client 148 self.protos = protos 149 150 # Store objects with references to RPC services, sorted by package. 151 self._services: Dict[str, types.SimpleNamespace] = defaultdict( 152 types.SimpleNamespace) 153 154 self._variables: Dict[str, object] = dict( 155 Status=pw_status.Status, 156 set_target=functions.help_as_repr(self.set_target), 157 # The original built-in help function is available as 'python_help'. 158 python_help=help, 159 ) 160 161 # Make the RPC clients and protos available in the console. 162 self._variables.update((c.name, c.client) for c in self.client_info) 163 164 # Make the proto package hierarchy directly available in the console. 165 for package in self.protos.packages: 166 self._variables[package._package] = package # pylint: disable=protected-access 167 168 # Monkey patch the message types to use an improved repr function. 169 for message_type in self.protos.messages(): 170 message_type.__repr__ = python_protos.proto_repr 171 172 # Set up the 'help' command. 173 all_methods = chain.from_iterable(c.rpc_client.methods() 174 for c in self.client_info) 175 self._helper = CommandHelper.from_methods( 176 all_methods, self._variables, help_header, 177 'Type a command and hit Enter to see detailed help information.') 178 179 self._variables['help'] = self._helper 180 181 # Call set_target to set up for the default target. 182 self.set_target(self.current_client) 183 184 def variables(self) -> Dict[str, Any]: 185 """Returns a mapping of names to variables for use in an RPC console.""" 186 return self._variables 187 188 def set_target(self, 189 selected_client: object, 190 channel_id: int = None) -> None: 191 """Sets the default target for commands.""" 192 # Make sure the variable is one of the client variables. 193 name = '' 194 rpc_client: Any = None 195 196 for name, client, rpc_client in self.client_info: 197 if selected_client is client: 198 print('CURRENT RPC TARGET:', name) 199 break 200 else: 201 raise ValueError('Supported targets :' + 202 ', '.join(c.name for c in self.client_info)) 203 204 # Update the RPC services to use the newly selected target. 205 for service_client in rpc_client.channel(channel_id).rpcs: 206 # Patch all method protos to use the improved __repr__ function too. 207 for method in (m.method for m in service_client): 208 method.request_type.__repr__ = python_protos.proto_repr 209 method.response_type.__repr__ = python_protos.proto_repr 210 211 service = service_client._service # pylint: disable=protected-access 212 setattr(self._services[service.package], service.name, 213 service_client) 214 215 # Add the RPC methods to their proto packages. 216 for package_name, rpcs in self._services.items(): 217 self.protos.packages[package_name]._add_item(rpcs) # pylint: disable=protected-access 218 219 self.current_client = selected_client 220