cannot_reference_dunder_param.py

import functools
import time
from collections import abc
from typing import Callable, Optional, Union

Step = abc.Callable


def step(description: str) -> Step:
    def _step(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        wrapper.__step_description__ = description
        return wrapper
    return _step


class TimeTracker:
    def __init__(self, printer: Optional[Callable] = None):
        timestamp = time.time()
        self.t_start_origin: float = timestamp
        self.t_start: float = timestamp
        self.time_table: dict = {}
        self.printer = printer or print

    def log_event(self, event: str):
        self.time_table[event] = time.time() - self.t_start
        self.t_start = time.time()

    def __call__(self, event: Union[str, Step]) -> None:
        if isinstance(event, str):
            return self.log_event(event)
        elif isinstance(event, Step):
            event()
            return self.log_event(event.__step_description__)
        else:
            raise NotImplementedError('Cannot log objects of type {type(event)}')

    def __enter__(self, *args, **kwargs):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.print_report()

    def print_report(self, printer: Optional[Callable] = None):
        printer = printer or self.printer
        printer('{: <20} {: <20}'.format('Message', 'Took'))
        for msg, took in self.time_table.items():
            printer('{: <70} {: <20.6f}'.format(msg, took))

        total = time.time() - self.t_start_origin
        printer('{: <20} {: <20.6f}'.format('Total', total))