algo_present/visualize.py

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))