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

28 statements  

« prev     ^ index     » next       coverage.py v6.4.4, created at 2023-07-17 14:22 -0600

1from __future__ import annotations 

2 

3import numpy as np 

4 

5from pandas._typing import NumpyIndexT 

6 

7from pandas.core.dtypes.common import is_list_like 

8 

9 

10def cartesian_product(X) -> list[np.ndarray]: 

11 """ 

12 Numpy version of itertools.product. 

13 Sometimes faster (for large inputs)... 

14 

15 Parameters 

16 ---------- 

17 X : list-like of list-likes 

18 

19 Returns 

20 ------- 

21 product : list of ndarrays 

22 

23 Examples 

24 -------- 

25 >>> cartesian_product([list('ABC'), [1, 2]]) 

26 [array(['A', 'A', 'B', 'B', 'C', 'C'], dtype='<U1'), array([1, 2, 1, 2, 1, 2])] 

27 

28 See Also 

29 -------- 

30 itertools.product : Cartesian product of input iterables. Equivalent to 

31 nested for-loops. 

32 """ 

33 msg = "Input must be a list-like of list-likes" 

34 if not is_list_like(X): 

35 raise TypeError(msg) 

36 for x in X: 

37 if not is_list_like(x): 

38 raise TypeError(msg) 

39 

40 if len(X) == 0: 

41 return [] 

42 

43 lenX = np.fromiter((len(x) for x in X), dtype=np.intp) 

44 cumprodX = np.cumproduct(lenX) 

45 

46 if np.any(cumprodX < 0): 

47 raise ValueError("Product space too large to allocate arrays!") 

48 

49 a = np.roll(cumprodX, 1) 

50 a[0] = 1 

51 

52 if cumprodX[-1] != 0: 

53 b = cumprodX[-1] / cumprodX 

54 else: 

55 # if any factor is empty, the cartesian product is empty 

56 b = np.zeros_like(cumprodX) 

57 

58 # error: Argument of type "int_" cannot be assigned to parameter "num" of 

59 # type "int" in function "tile_compat" 

60 return [ 

61 tile_compat( 

62 np.repeat(x, b[i]), 

63 np.product(a[i]), # pyright: ignore[reportGeneralTypeIssues] 

64 ) 

65 for i, x in enumerate(X) 

66 ] 

67 

68 

69def tile_compat(arr: NumpyIndexT, num: int) -> NumpyIndexT: 

70 """ 

71 Index compat for np.tile. 

72 

73 Notes 

74 ----- 

75 Does not support multi-dimensional `num`. 

76 """ 

77 if isinstance(arr, np.ndarray): 

78 return np.tile(arr, num) 

79 

80 # Otherwise we have an Index 

81 taker = np.tile(np.arange(len(arr)), num) 

82 return arr.take(taker)