"""
Performance utilities.
"""

import gc
from contextlib import contextmanager
from time import perf_counter

from .. import logger

logger = logger.new_logger(__name__)


@contextmanager
def timer(identifier: str):
    """Utility context manager for timing code snippets.

    Args:
        identifier (str): Identifier for the code snippet, will appear in the
        log message.

    Example:
        with timer("my_code"):
            pass
    """
    start = perf_counter()
    yield
    end = perf_counter()
    elapsed = end - start
    logger.info("{} executed in {} sec.", identifier, round(elapsed, 3))


@contextmanager
def no_garbage_collection(collector=gc):
    """Context manager to disable garbage collection. This is useful for performance
    testing and profiling, as garbage collection can skew results. It is also useful
    to speed up code execution in some cases.

    Args:
        collector (gc.Collector): Garbage collector to use. Defaults to the built-in
            garbage collector module. This is useful for providing a mock object for
            testing and not having to mock the built-in module.
    """
    was_enabled = collector.isenabled()
    if was_enabled:
        logger.debug("Disabling garbage collection...")
    collector.disable()
    try:
        yield
    finally:
        if was_enabled:
            logger.debug("Re-enabling garbage collection...")
            collector.enable()