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 __future__ import annotations 17 18from collections import defaultdict 19import functools 20from itertools import chain 21import inspect 22import textwrap 23import types 24from typing import ( 25 Any, 26 Collection, 27 Iterable, 28 Mapping, 29 NamedTuple, 30) 31 32import pw_status 33from pw_protobuf_compiler import python_protos 34 35import pw_rpc 36from pw_rpc.descriptors import Method 37from pw_rpc.console_tools import functions 38 39_INDENT = ' ' 40 41 42class CommandHelper: 43 """Used to implement a help command in an RPC console.""" 44 45 @classmethod 46 def from_methods( 47 cls, 48 methods: Iterable[Method], 49 variables: Mapping[str, object], 50 header: str, 51 footer: str = '', 52 ) -> CommandHelper: 53 return cls({m.full_name: m for m in methods}, variables, header, footer) 54 55 def __init__( 56 self, 57 methods: Mapping[str, object], 58 variables: Mapping[str, object], 59 header: str, 60 footer: str = '', 61 ): 62 self._methods = methods 63 self._variables = variables 64 self.header = header 65 self.footer = footer 66 67 def help(self, item: object = None) -> str: 68 """Returns a help string with a command or all commands listed.""" 69 70 if item is None: 71 all_vars = '\n'.join(sorted(self._variables_without_methods())) 72 all_rpcs = '\n'.join(self._methods) 73 return ( 74 f'{self.header}\n\n' 75 f'All variables:\n\n{textwrap.indent(all_vars, _INDENT)}' 76 '\n\n' 77 f'All commands:\n\n{textwrap.indent(all_rpcs, _INDENT)}' 78 f'\n\n{self.footer}'.strip() 79 ) 80 81 # If item is a string, find commands matching that. 82 if isinstance(item, str): 83 matches = {n: m for n, m in self._methods.items() if item in n} 84 if not matches: 85 return f'No matches found for {item!r}' 86 87 if len(matches) == 1: 88 name, method = next(iter(matches.items())) 89 return f'{name}\n\n{inspect.getdoc(method)}' 90 91 return f'Multiple matches for {item!r}:\n\n' + textwrap.indent( 92 '\n'.join(matches), _INDENT 93 ) 94 95 return inspect.getdoc(item) or f'No documentation for {item!r}.' 96 97 def _variables_without_methods(self) -> Mapping[str, object]: 98 packages = frozenset( 99 n.split('.', 1)[0] for n in self._methods if '.' in n 100 ) 101 102 return { 103 name: var 104 for name, var in self._variables.items() 105 if name not in packages 106 } 107 108 def __call__(self, item: object = None) -> None: 109 """Prints the help string.""" 110 print(self.help(item)) 111 112 def __repr__(self) -> str: 113 """Returns the help, so foo and foo() are equivalent in a console.""" 114 return self.help() 115 116 117class ClientInfo(NamedTuple): 118 """Information about an RPC client as it appears in the console.""" 119 120 # The name to use in the console to refer to this client. 121 name: str 122 123 # An object to use in the console for the client. May be a pw_rpc.Client. 124 client: object 125 126 # The pw_rpc.Client; may be the same object as client. 127 rpc_client: pw_rpc.Client 128 129 130def flattened_rpc_completions( 131 client_info_list: Collection[ClientInfo], 132) -> dict[str, str]: 133 """Create a flattened list of rpc commands for repl auto-completion. 134 135 This gathers all rpc commands from a set of ClientInfo variables and 136 produces a flattened list of valid rpc commands to run in an RPC 137 console. This is useful for passing into 138 prompt_toolkit.completion.WordCompleter. 139 140 Args: 141 client_info_list: List of ClientInfo variables 142 143 Returns: 144 Dict of flattened rpc commands as keys, and 'RPC' as values. 145 For example: :: 146 147 { 148 'device.rpcs.pw.rpc.EchoService.Echo': 'RPC, 149 'device.rpcs.pw.rpc.BatteryService.GetBatteryStatus': 'RPC', 150 } 151 """ 152 rpc_list = list( 153 chain.from_iterable( 154 [ 155 '{}.rpcs.{}'.format(c.name, a.full_name) 156 for a in c.rpc_client.methods() 157 ] 158 for c in client_info_list 159 ) 160 ) 161 162 # Dict should contain completion text as keys and descriptions as values. 163 custom_word_completions = { 164 flattened_rpc_name: 'RPC' for flattened_rpc_name in rpc_list 165 } 166 return custom_word_completions 167 168 169class Context: 170 """The Context class is used to set up an interactive RPC console. 171 172 The Context manages a set of variables that make it easy to access RPCs and 173 protobufs in a REPL. 174 175 As an example, this class can be used to set up a console with IPython: 176 177 .. code-block:: python 178 179 context = console_tools.Context( 180 clients, default_client, protos, help_header=WELCOME_MESSAGE) 181 IPython.start_ipython(argv=[], user_ns=dict(**context.variables())) 182 """ 183 184 def __init__( 185 self, 186 client_info: Collection[ClientInfo], 187 default_client: Any, 188 protos: python_protos.Library, 189 *, 190 help_header: str = '', 191 ) -> None: 192 """Creates an RPC console context. 193 194 Protos and RPC services are accessible by their proto package and name. 195 The target for these can be set with the set_target function. 196 197 Args: 198 client_info: ClientInfo objects that represent the clients this 199 console uses to communicate with other devices 200 default_client: default client object; must be one of the clients 201 protos: protobufs to use for RPCs for all clients 202 help_header: Message to display for the help command 203 """ 204 assert client_info, 'At least one client must be provided!' 205 206 self.client_info = client_info 207 self.current_client = default_client 208 self.protos = protos 209 210 # Store objects with references to RPC services, sorted by package. 211 self._services: dict[str, types.SimpleNamespace] = defaultdict( 212 types.SimpleNamespace 213 ) 214 215 self._variables: dict[str, object] = dict( 216 Status=pw_status.Status, 217 set_target=functions.help_as_repr(self.set_target), 218 # The original built-in help function is available as 'python_help'. 219 python_help=help, 220 ) 221 222 # Make the RPC clients and protos available in the console. 223 self._variables.update((c.name, c.client) for c in self.client_info) 224 225 # Make the proto package hierarchy directly available in the console. 226 for package in self.protos.packages: 227 self._variables[ 228 package._package 229 ] = package # pylint: disable=protected-access 230 231 # Monkey patch the message types to use an improved repr function. 232 for message_type in self.protos.messages(): 233 message_type.__repr__ = python_protos.proto_repr 234 235 # Set up the 'help' command. 236 all_methods = chain.from_iterable( 237 c.rpc_client.methods() for c in self.client_info 238 ) 239 self._helper = CommandHelper.from_methods( 240 all_methods, 241 self._variables, 242 help_header, 243 'Type a command and hit Enter to see detailed help information.', 244 ) 245 246 self._variables['help'] = self._helper 247 248 # Call set_target to set up for the default target. 249 self.set_target(self.current_client) 250 251 def flattened_rpc_completions(self): 252 """Create a flattened list of rpc commands for repl auto-completion.""" 253 return flattened_rpc_completions(self.client_info) 254 255 def variables(self) -> dict[str, Any]: 256 """Returns a mapping of names to variables for use in an RPC console.""" 257 return self._variables 258 259 def set_target( 260 self, selected_client: object, channel_id: int | None = None 261 ) -> None: 262 """Sets the default target for commands.""" 263 # Make sure the variable is one of the client variables. 264 name = '' 265 rpc_client: Any = None 266 267 for name, client, rpc_client in self.client_info: 268 if selected_client is client: 269 print('CURRENT RPC TARGET:', name) 270 break 271 else: 272 raise ValueError( 273 'Supported targets :' 274 + ', '.join(c.name for c in self.client_info) 275 ) 276 277 # Update the RPC services to use the newly selected target. 278 for service_client in rpc_client.channel(channel_id).rpcs: 279 # Patch all method protos to use the improved __repr__ function too. 280 for method in (m.method for m in service_client): 281 method.request_type.__repr__ = python_protos.proto_repr 282 method.response_type.__repr__ = python_protos.proto_repr 283 284 service = ( 285 service_client._service # pylint: disable=protected-access 286 ) 287 setattr( 288 self._services[service.package], service.name, service_client 289 ) 290 291 # Add the RPC methods to their proto packages. 292 for package_name, rpcs in self._services.items(): 293 # pylint: disable=protected-access 294 self.protos.packages[package_name]._add_item(rpcs) 295 # pylint: enable=protected-access 296 297 self.current_client = selected_client 298 299 300def _create_command_alias(command: Any, name: str, message: str) -> object: 301 """Wraps __call__, __getattr__, and __repr__ to print a message.""" 302 303 @functools.wraps(command.__call__) 304 def print_message_and_call(_, *args, **kwargs): 305 print(message) 306 return command(*args, **kwargs) 307 308 def getattr_and_print_message(_, name: str) -> Any: 309 attr = getattr(command, name) 310 print(message) 311 return attr 312 313 return type( 314 name, 315 (), 316 dict( 317 __call__=print_message_and_call, 318 __getattr__=getattr_and_print_message, 319 __repr__=lambda _: message, 320 ), 321 )() 322 323 324def _access_in_dict_or_namespace(item, name: str, create_if_missing: bool): 325 """Gets name as either a key or attribute on item.""" 326 try: 327 return item[name] 328 except KeyError: 329 if create_if_missing: 330 try: 331 item[name] = types.SimpleNamespace() 332 return item[name] 333 except TypeError: 334 pass 335 except TypeError: 336 pass 337 338 if create_if_missing and not hasattr(item, name): 339 setattr(item, name, types.SimpleNamespace()) 340 341 return getattr(item, name) 342 343 344def _access_names(item, names: Iterable[str], create_if_missing: bool): 345 for name in names: 346 item = _access_in_dict_or_namespace(item, name, create_if_missing) 347 348 return item 349 350 351def alias_deprecated_command( 352 variables: Any, old_name: str, new_name: str 353) -> None: 354 """Adds an alias for an old command that redirects to the new command. 355 356 The deprecated command prints a message then invokes the new command. 357 """ 358 # Get the new command. 359 item = _access_names( 360 variables, new_name.split('.'), create_if_missing=False 361 ) 362 363 # Create a wrapper to the new comamnd with the old name. 364 wrapper = _create_command_alias( 365 item, 366 old_name, 367 f'WARNING: {old_name} is DEPRECATED; use {new_name} instead', 368 ) 369 370 # Add the wrapper to the variables with the old command's name. 371 name_parts = old_name.split('.') 372 item = _access_names(variables, name_parts[:-1], create_if_missing=True) 373 setattr(item, name_parts[-1], wrapper) 374