Source code for ewoksid14.timepix4.t4_profile

import functools
import os
import time
import tracemalloc
from contextlib import ExitStack
from contextlib import contextmanager
from datetime import timedelta
from typing import List
from typing import Union

import numpy
from memory_profiler import profile as mem_line_profile  # noqa F401




def main_profile(mem=True, time=True):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            with ExitStack() as stack:
                if mem:
                    stack.enter_context(mem_context())
                if time:
                    stack.enter_context(time_context())
                return func(*args, **kwargs)

        return wrapper

    return decorator





@contextmanager
def mem_context():
    tracemalloc.start()
    try:
        yield
    finally:
        current, peak = tracemalloc.get_traced_memory()
        print(f"Current memory usage: {_format_bytes(current)}")
        print(f"Peak memory usage: {_format_bytes(peak)}")





@contextmanager
def time_context():
    t0 = time.time()
    try:
        yield
    finally:
        duration = int(time.time() - t0)
        print(f"\nDuration: {timedelta(seconds=duration)}")





def print_file_size(files: Union[List[str], List[List[str]]]) -> None:
    if not isinstance(files[0], str):
        files = [f for sublist in files for f in sublist]
    total_bytes = sum(os.path.getsize(f) for f in files)
    print(f"Total file size: {_format_bytes(total_bytes)}")





def print_ndarray_stats(data: numpy.ndarray) -> None:
    total_bytes = data.nbytes
    print(f"Total ndarray size: {_format_bytes(total_bytes)}")
    print(f" Dtype: {data.dtype}")
    print(f" Shape: {data.shape}")

    if data.size == 0:
        return

    print(f" Min: {data.min():,}")
    print(f" Max: {data.max():,}")
    print(f" Sum: {data.sum():,}")



def _format_bytes(num_bytes: int) -> str:
    """
    Convert a byte value into a human-readable string with
    automatically selected unit.
    """
    units = ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]
    value = float(num_bytes)

    for unit in units:
        if value < 1024 or unit == units[-1]:
            return f"{value:.2f} {unit}"
        value /= 1024