diff --git a/visualize.py b/visualize.py new file mode 100644 index 0000000..645d67d --- /dev/null +++ b/visualize.py @@ -0,0 +1,39 @@ +import functools +import os +from typing import Any, Callable + +""" +This file contains some decorators for visualization. +To use all of this decorators, you MUST set env RUNNING_UNITTESTS=1. +""" + + +def show_name_args(func: Callable) -> Callable: + """ + A decorator that prints function information when running under unittests. + + Args: + func (Callable): The function to wrap with the decorator. + + Returns: + Callable: The wrapped function with conditional logging. + """ + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + # Check if running under unittest environment (e.g., an environment variable is set) + if os.getenv("RUNNING_UNITTESTS"): + print(f"Calling {func.__name__} with args {args} and kwargs {kwargs}") + return func(*args, **kwargs) + + return wrapper + + +if __name__ == "__main__": + # Example of using the decorator + @show_name_args + def my_function(x: int, y: int) -> int: + """Add two numbers.""" + return x + y + + print(my_function(1, 4))