Coverage for /var/srv/projects/api.amasfac.comuna18.com/tmp/venv/lib/python3.9/site-packages/pandas/core/util/numba_.py: 22%

1"""Common utilities for Numba operations""" 

2from __future__ import annotations 

3 

4import types 

5from typing import ( 

6 TYPE_CHECKING, 

7 Callable, 

8) 

9 

10import numpy as np 

11 

12from pandas.compat._optional import import_optional_dependency 

13from pandas.errors import NumbaUtilError 

14 

15GLOBAL_USE_NUMBA: bool = False 

16 

17 

18def maybe_use_numba(engine: str | None) -> bool: 

19 """Signal whether to use numba routines.""" 

20 return engine == "numba" or (engine is None and GLOBAL_USE_NUMBA) 

21 

22 

23def set_use_numba(enable: bool = False) -> None: 

24 global GLOBAL_USE_NUMBA 

25 if enable: 

26 import_optional_dependency("numba") 

27 GLOBAL_USE_NUMBA = enable 

28 

29 

30def get_jit_arguments( 

31 engine_kwargs: dict[str, bool] | None = None, kwargs: dict | None = None 

32) -> dict[str, bool]: 

33 """ 

34 Return arguments to pass to numba.JIT, falling back on pandas default JIT settings. 

35 

36 Parameters 

37 ---------- 

38 engine_kwargs : dict, default None 

39 user passed keyword arguments for numba.JIT 

40 kwargs : dict, default None 

41 user passed keyword arguments to pass into the JITed function 

42 

43 Returns 

44 ------- 

45 dict[str, bool] 

46 nopython, nogil, parallel 

47 

48 Raises 

49 ------ 

50 NumbaUtilError 

51 """ 

52 if engine_kwargs is None: 

53 engine_kwargs = {} 

54 

55 nopython = engine_kwargs.get("nopython", True) 

56 if kwargs and nopython: 

57 raise NumbaUtilError( 

58 "numba does not support kwargs with nopython=True: " 

59 "https://github.com/numba/numba/issues/2916" 

60 ) 

61 nogil = engine_kwargs.get("nogil", False) 

62 parallel = engine_kwargs.get("parallel", False) 

63 return {"nopython": nopython, "nogil": nogil, "parallel": parallel} 

64 

65 

66def jit_user_function( 

67 func: Callable, nopython: bool, nogil: bool, parallel: bool 

68) -> Callable: 

69 """ 

70 JIT the user's function given the configurable arguments. 

71 

72 Parameters 

73 ---------- 

74 func : function 

75 user defined function 

76 nopython : bool 

77 nopython parameter for numba.JIT 

78 nogil : bool 

79 nogil parameter for numba.JIT 

80 parallel : bool 

81 parallel parameter for numba.JIT 

82 

83 Returns 

84 ------- 

85 function 

86 Numba JITed function 

87 """ 

88 if TYPE_CHECKING: 

89 import numba 

90 else: 

91 numba = import_optional_dependency("numba") 

92 

93 if numba.extending.is_jitted(func): 

94 # Don't jit a user passed jitted function 

95 numba_func = func 

96 else: 

97 

98 @numba.generated_jit(nopython=nopython, nogil=nogil, parallel=parallel) 

99 def numba_func(data, *_args): 

100 if getattr(np, func.__name__, False) is func or isinstance( 

101 func, types.BuiltinFunctionType 

102 ): 

103 jf = func 

104 else: 

105 jf = numba.jit(func, nopython=nopython, nogil=nogil) 

106 

107 def impl(data, *_args): 

108 return jf(data, *_args) 

109 

110 return impl 

111 

112 return numba_func