helpers ¶

def is_from_vs_module(obj: Any) -> bool:
    """Returns true if the module in which the obj was defined is VapourSynth."""

    global _vs_module

    if not _vs_module:
        import vapoursynth

        _vs_module = vapoursynth

    if hasattr(obj, "__module__"):
        return sys.modules[obj.__module__] is _vs_module

    return getmodule(type(obj)) is _vs_module

def lazy_load(name: str, package: str | None = None, exc: Callable[[], Exception] | None = None) -> ModuleType:
    """Lazily load a package."""

    if name in sys.modules:
        return sys.modules[name]

    from importlib.util import LazyLoader, find_spec, module_from_spec

    spec = find_spec(name, package)

    if spec is None:
        raise exc() if exc else ModuleNotFoundError(f"No module named {name!r}", name=name, path=__file__)

    module = module_from_spec(spec)

    if spec.loader is None:
        raise exc() if exc else NotImplementedError

    loader = LazyLoader(spec.loader)
    loader.exec_module(module)

    return module

def require_jet_dependency[**P, R](
    *name: Literal["scipy", "rich", "psutil"],
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    """Decorator that raises DependencyNotFoundError when a specific package is missing."""

    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            try:
                return func(*args, **kwargs)
            except ImportError as e:
                if e.name in name:
                    exc = DependencyNotFoundError(
                        func,
                        e.name,
                        "Missing dependency '{package}' for function '{func_name}'. Please install vsjetpack[full]",
                        func_name=func,
                    )

                    raise exc from None
            raise

        return wrapper

    return decorator

@require_jet_dependency("rich")
def setup_logging(level: str | int = INFO, **kwargs: Any) -> None:
    """
    Configure global logging.

    Args:
        level: Log level. Defaults to INFO.
        kwargs: Arguments forwarded to logging.basicConfig
    """
    from rich.console import Console
    from rich.logging import RichHandler
    from rich.text import Text

    class CustomJetHandler(RichHandler):
        def format(self, record: LogRecord) -> str:
            # Return a modified shallow copy of the LogRecord with transformed
            # parameters for specific loggers.
            record = copy(record)
            if isinstance(record.args, tuple):
                transformed = _transform_record_args(dict(enumerate(record.args)))
                record.args = tuple(transformed.values())
            elif isinstance(record.args, Mapping):
                record.args = _transform_record_args(record.args)

            return super().format(record)

    handler = CustomJetHandler(
        console=Console(stderr=True),
        omit_repeated_times=False,
        show_time=True,
        rich_tracebacks=True,
        log_time_format=lambda dt: Text("[{}.{:03d}]".format(dt.strftime("%H:%M:%S"), dt.microsecond // 1000)),
    )
    handler.setFormatter(Formatter("{name}: {message}", style="{"))

    for module in _JETPACK_MODULES:
        logger = getLogger(module)
        logger.setLevel(level)
        logger.propagate = False

        if not logger.hasHandlers():
            logger.addHandler(handler)

    basicConfig(level=level, **kwargs)