observer.py

# Copyright (c) OpenMMLab. All rights reserved.
from typing import Any, Callable


class Observer:
    """The Observer class applies a user-specified function on its inputs and
    stores the results in a buffer.

    Args:
        observe_fn (Callable[..., Any]): The function to apply on inputs.
    """

    def __init__(self, observe_fn: Callable[..., Any]) -> None:
        super().__init__()
        self.fn = observe_fn
        self.buffer = list()
        self.enabled = False

    def enable_observer(self, enabled: bool = True) -> None:
        """Enable or disable the observer.

        Args:
            enabled (bool, optional): Whether to enable the observer.
                Defaults to True.
        """
        self.enabled = enabled

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """Apply the observer function on the input if the observer is enabled.
        Args:
            *args: Variable length argument list.
            **kwargs: Arbitrary keyword arguments.
        """
        if self.enabled:
            self.buffer.append(self.fn(*args, **kwargs))