Coverage for src/evutils/_jit.py: 100%

1"""Lazy numba compilation helper. 

2 

3Importing numba costs several hundred milliseconds, which would be paid by 

4``import evutils.io`` even when no numba-accelerated code path is ever used 

5(e.g. pure reading through the C parsers). :func:`lazy_njit` defers both the 

6numba import and the JIT compilation to the first call of the decorated 

7function. 

8""" 

9from __future__ import annotations 

10 

11import functools 

12from typing import Any, Callable, TypeVar 

13 

14F = TypeVar("F", bound=Callable[..., Any]) 

15 

16 

17def lazy_njit(fn: F) -> F: 

18 """``numba.njit``, but imported and compiled on first call. 

19 

20 The wrapped function behaves like the ``@nb.njit``-decorated original; 

21 only the timing of the numba import/compilation differs. The per-call 

22 overhead after the first call is a single ``is None`` check. 

23 """ 

24 compiled: Callable[..., Any] | None = None 

25 

26 @functools.wraps(fn) 

27 def wrapper(*args: Any, **kwargs: Any) -> Any: 

28 nonlocal compiled 

29 if compiled is None: 

30 import numba as nb 

31 compiled = nb.njit(fn) 

32 return compiled(*args, **kwargs) 

33 

34 return wrapper # type: ignore[return-value]