xref: /aosp_15_r20/external/pigweed/pw_rpc/py/pw_rpc/console_tools/functions.py (revision 61c4878ac05f98d0ceed94b57d316916de578985) 
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"""Code for improving interactive use of Python functions."""
15
16from __future__ import annotations
17
18import inspect
19import textwrap
20from typing import Callable
21
22
23def _annotation_name(annotation: object) -> str:
24    if isinstance(annotation, str):
25        return annotation
26
27    return getattr(annotation, '__name__', repr(annotation))
28
29
30def format_parameter(param: inspect.Parameter) -> str:
31    """Formats a parameter for printing in a function signature."""
32    if param.kind == param.VAR_POSITIONAL:
33        name = '*' + param.name
34    elif param.kind == param.VAR_KEYWORD:
35        name = '**' + param.name
36    else:
37        name = param.name
38
39    if param.default is param.empty:
40        default = ''
41    else:
42        default = f' = {param.default}'
43
44    if param.annotation is param.empty:
45        annotation = ''
46    else:
47        annotation = f': {_annotation_name(param.annotation)}'
48
49    return f'{name}{annotation}{default}'
50
51
52def format_signature(name: str, signature: inspect.Signature) -> str:
53    """Formats a function signature as if it were source code.
54
55    Does not yet handle / and * markers.
56    """
57    params = ', '.join(
58        format_parameter(arg) for arg in signature.parameters.values()
59    )
60    if signature.return_annotation is signature.empty:
61        return_annotation = ''
62    else:
63        return_annotation = ' -> ' + _annotation_name(
64            signature.return_annotation
65        )
66
67    return f'{name}({params}){return_annotation}'
68
69
70def format_function_help(function: Callable) -> str:
71    """Formats a help string with a declaration and docstring."""
72    signature = format_signature(
73        function.__name__, inspect.signature(function, follow_wrapped=False)
74    )
75
76    docs = inspect.getdoc(function) or '(no docstring)'
77    return f'{signature}:\n\n{textwrap.indent(docs, "    ")}'
78
79
80def help_as_repr(function: Callable) -> Callable:
81    """Wraps a function so that its repr() and docstring provide detailed help.
82
83    This is useful for creating commands in an interactive console. In a
84    console, typing a function's name and hitting Enter shows rich documentation
85    with the full function signature, type annotations, and docstring when the
86    function is wrapped with help_as_repr.
87    """
88
89    def display_help(_):
90        return format_function_help(function)
91
92    return type(
93        function.__name__,
94        (),
95        dict(
96            __call__=staticmethod(function),
97            __doc__=format_function_help(function),
98            __repr__=display_help,
99        ),
100    )()
101