1# echo.py: Tracing function calls using Python decorators.
2#
3# Written by Thomas Guest <[email protected]>
4# Please see http://wordaligned.org/articles/echo
5#
6# Place into the public domain.
7
8"""Echo calls made to functions and methods in a module.
9
10"Echoing" a function call means printing out the name of the function
11and the values of its arguments before making the call (which is more
12commonly referred to as "tracing", but Python already has a trace module).
13
14Example: to echo calls made to functions in "my_module" do:
15
16  import echo
17  import my_module
18  echo.echo_module(my_module)
19
20Example: to echo calls made to functions in "my_module.my_class" do:
21
22  echo.echo_class(my_module.my_class)
23
24Alternatively, echo.echo can be used to decorate functions. Calls to the
25decorated function will be echoed.
26
27Example:
28-------
29  @echo.echo
30  def my_function(args):
31      pass
32
33
34"""
35
36from __future__ import annotations
37
38import inspect
39import sys
40from typing import TYPE_CHECKING
41
42if TYPE_CHECKING:
43    from types import MethodType
44    from typing import Any, Callable
45
46
47def name(item: Callable) -> str:
48    """Return an item's name."""
49    return item.__name__
50
51
52def is_classmethod(instancemethod: MethodType, klass: type) -> bool:
53    """Determine if an instancemethod is a classmethod."""
54    return inspect.ismethod(instancemethod) and instancemethod.__self__ is klass
55
56
57def is_static_method(method: MethodType, klass: type) -> bool:
58    """Returns True if method is an instance method of klass."""
59    return next(
60        (isinstance(c.__dict__[name(method)], staticmethod) for c in klass.mro() if name(method) in c.__dict__),
61        False,
62    )
63
64
65def is_class_private_name(name: str) -> bool:
66    """Determine if a name is a class private name."""
67    # Exclude system defined names such as __init__, __add__ etc
68    return name.startswith("__") and not name.endswith("__")
69
70
71def method_name(method: MethodType) -> str:
72    """Return a method's name.
73
74    This function returns the name the method is accessed by from
75    outside the class (i.e. it prefixes "private" methods appropriately).
76    """
77    mname = name(method)
78    if is_class_private_name(mname):
79        mname = f"_{name(method.__self__.__class__)}{mname}"
80    return mname
81
82
83def format_arg_value(arg_val: tuple[str, tuple[Any, ...]]) -> str:
84    """Return a string representing a (name, value) pair."""
85    arg, val = arg_val
86    return f"{arg}={val!r}"
87
88
89def echo(fn: Callable, write: Callable[[str], int | None] = sys.stdout.write) -> Callable:
90    """Echo calls to a function.
91
92    Returns a decorated version of the input function which "echoes" calls
93    made to it by writing out the function's name and the arguments it was
94    called with.
95    """
96    import functools
97
98    # Unpack function's arg count, arg names, arg defaults
99    code = fn.__code__
100    argcount = code.co_argcount
101    argnames = code.co_varnames[:argcount]
102    fn_defaults: tuple[Any] = fn.__defaults__ or ()
103    argdefs = dict(list(zip(argnames[-len(fn_defaults) :], fn_defaults)))
104
105    @functools.wraps(fn)
106    def wrapped(*v: Any, **k: Any) -> Callable:
107        # Collect function arguments by chaining together positional,
108        # defaulted, extra positional and keyword arguments.
109        positional = list(map(format_arg_value, list(zip(argnames, v))))
110        defaulted = [format_arg_value((a, argdefs[a])) for a in argnames[len(v) :] if a not in k]
111        nameless = list(map(repr, v[argcount:]))
112        keyword = list(map(format_arg_value, list(k.items())))
113        args = positional + defaulted + nameless + keyword
114        write(f"{name(fn)}({', '.join(args)})\n")
115        return fn(*v, **k)
116
117    return wrapped
118
119
120def echo_instancemethod(klass: type, method: MethodType, write: Callable[[str], int | None] = sys.stdout.write) -> None:
121    """Change an instancemethod so that calls to it are echoed.
122
123    Replacing a classmethod is a little more tricky.
124    See: http://www.python.org/doc/current/ref/types.html
125    """
126    mname = method_name(method)
127
128    # Avoid recursion printing method calls
129    if mname in {"__str__", "__repr__"}:
130        return
131
132    if is_classmethod(method, klass):
133        setattr(klass, mname, classmethod(echo(method.__func__, write)))
134    else:
135        setattr(klass, mname, echo(method, write))
136
137
138def echo_class(klass: type, write: Callable[[str], int | None] = sys.stdout.write) -> None:
139    """Echo calls to class methods and static functions"""
140    for _, method in inspect.getmembers(klass, inspect.ismethod):
141        # In python 3 only class methods are returned here
142        echo_instancemethod(klass, method, write)
143    for _, fn in inspect.getmembers(klass, inspect.isfunction):
144        if is_static_method(fn, klass):
145            setattr(klass, name(fn), staticmethod(echo(fn, write)))
146        else:
147            # It's not a class or a static method, so it must be an instance method.
148            echo_instancemethod(klass, fn, write)
149
150
151def echo_module(mod: MethodType, write: Callable[[str], int | None] = sys.stdout.write) -> None:
152    """Echo calls to functions and methods in a module."""
153    for fname, fn in inspect.getmembers(mod, inspect.isfunction):
154        setattr(mod, fname, echo(fn, write))
155    for _, klass in inspect.getmembers(mod, inspect.isclass):
156        echo_class(klass, write)
157