GeometryForcing/utils/print_utils.py

import warnings
import os
import sys
from functools import wraps
from contextlib import contextmanager
from colorama import Fore


def cyan(x: str) -> str:
    return f"{Fore.CYAN}{x}{Fore.RESET}"


@contextmanager
def suppress_print():
    with open(os.devnull, "w") as devnull:
        old_stdout = sys.stdout
        sys.stdout = devnull
        try:
            yield
        finally:
            sys.stdout = old_stdout


class suppress_warnings:
    """
    A context manager to suppress warnings. This can be used as a decorator as well.
    Example:
    ```
    @suppress_warnings()
    def my_function():
        pass

    with suppress_warnings():
        another_function()
    ```
    """

    def __init__(self, category=Warning):
        self.category = category

    def __call__(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            with self:
                return func(*args, **kwargs)

        return wrapper

    def __enter__(self):
        self.warnings_manager = warnings.catch_warnings()
        self.warnings_manager.__enter__()
        warnings.simplefilter("ignore", self.category)

    def __exit__(self, exc_type, exc_value, traceback):
        self.warnings_manager.__exit__(exc_type, exc_value, traceback)


def once_per_key():
    keys = set()

    def decorator(fn):
        def wrapper(*args, **kwargs):
            key = kwargs.get("key", args[0])
            if key not in keys:
                keys.add(key)
                fn(*args, **kwargs)

        return wrapper

    return decorator


@once_per_key()
def print_once(*args, **kwargs):
    print(*args, **kwargs)