import sys
import time
from importlib import machinery


class CustomLoader(machinery.SourceFileLoader):
    last_end = None
    depth = 0
    print_depth = None

    def exec_module(self, module):
        start = time.time_ns()
        CustomLoader.depth += 1
        try:
            result = super().exec_module(module)
        finally:
            CustomLoader.depth -= 1
            end = time.time_ns()
            if CustomLoader.last_end is None:
                CustomLoader.last_end = end
            if (
                CustomLoader.print_depth is None
                or CustomLoader.depth <= CustomLoader.print_depth
            ):
                print(
                    f"{' ' * CustomLoader.depth}{(end - start) // 10000 / 100}ms ({(end - CustomLoader.last_end) // 10000 / 100}ms local) to load {module.__name__} ",
                )
            CustomLoader.last_end = end
        return result


class CustomFinder(machinery.PathFinder):
    @classmethod
    def _get_spec(cls, fullname, path=None, target=None):
        spec = super()._get_spec(fullname, path, target)
        if spec is None:
            return None
        if isinstance(spec.loader, machinery.SourceFileLoader):
            spec.loader = CustomLoader(spec.loader.name, spec.loader.path)
        return spec


sys.meta_path.insert(4, CustomFinder())