Coverage for muutils/misc/numerical.py: 96%

1from __future__ import annotations 

2 

3 

4_SHORTEN_MAP: dict[int | float, str] = { 

5 1e3: "K", 

6 1e6: "M", 

7 1e9: "B", 

8 1e12: "t", 

9 1e15: "q", 

10 1e18: "Q", 

11} 

12 

13_SHORTEN_TUPLES: list[tuple[int | float, str]] = sorted( 

14 ((val, suffix) for val, suffix in _SHORTEN_MAP.items()), 

15 key=lambda x: -x[0], 

16) 

17 

18 

19_REVERSE_SHORTEN_MAP: dict[str, int | float] = {v: k for k, v in _SHORTEN_MAP.items()} 

20 

21 

22def shorten_numerical_to_str( 

23 num: int | float, 

24 small_as_decimal: bool = True, 

25 precision: int = 1, 

26) -> str: 

27 """shorten a large numerical value to a string 

28 1234 -> 1K 

29 

30 precision guaranteed to 1 in 10, but can be higher. reverse of `str_to_numeric` 

31 """ 

32 

33 # small values are returned as is 

34 num_abs: float = abs(num) 

35 if num_abs < 1e3: 

36 return str(num) 

37 

38 # iterate over suffixes from largest to smallest 

39 for i, (val, suffix) in enumerate(_SHORTEN_TUPLES): 

40 if num_abs > val or i == len(_SHORTEN_TUPLES) - 1: 

41 if (num_abs < val * 10) and small_as_decimal: 

42 return f"{num / val:.{precision}f}{suffix}" 

43 elif num_abs < val * 1e3: 

44 return f"{int(round(num / val))}{suffix}" 

45 

46 return f"{num:.{precision}f}" 

47 

48 

49def str_to_numeric( 

50 quantity: str, 

51 mapping: None | bool | dict[str, int | float] = True, 

52) -> int | float: 

53 """Convert a string representing a quantity to a numeric value. 

54 

55 The string can represent an integer, python float, fraction, or shortened via `shorten_numerical_to_str`. 

56 

57 # Examples: 

58 ``` 

59 >>> str_to_numeric("5") 

60 5 

61 >>> str_to_numeric("0.1") 

62 0.1 

63 >>> str_to_numeric("1/5") 

64 0.2 

65 >>> str_to_numeric("-1K") 

66 -1000.0 

67 >>> str_to_numeric("1.5M") 

68 1500000.0 

69 >>> str_to_numeric("1.2e2") 

70 120.0 

71 ``` 

72 

73 """ 

74 

75 # check is string 

76 if not isinstance(quantity, str): 

77 raise TypeError( 

78 f"quantity must be a string, got '{type(quantity) = }' '{quantity = }'" 

79 ) 

80 

81 # basic int conversion 

82 try: 

83 quantity_int: int = int(quantity) 

84 return quantity_int 

85 except ValueError: 

86 pass 

87 

88 # basic float conversion 

89 try: 

90 quantity_float: float = float(quantity) 

91 return quantity_float 

92 except ValueError: 

93 pass 

94 

95 # mapping 

96 _mapping: dict[str, int | float] 

97 if mapping is True or mapping is None: 

98 _mapping = _REVERSE_SHORTEN_MAP 

99 else: 

100 _mapping = mapping # type: ignore[assignment] 

101 

102 quantity_original: str = quantity 

103 

104 quantity = quantity.strip() 

105 

106 result: int | float 

107 multiplier: int | float = 1 

108 

109 # detect if it has a suffix 

110 suffixes_detected: list[bool] = [suffix in quantity for suffix in _mapping] 

111 n_suffixes_detected: int = sum(suffixes_detected) 

112 if n_suffixes_detected == 0: 

113 # no suffix 

114 pass 

115 elif n_suffixes_detected == 1: 

116 # find multiplier 

117 for suffix, mult in _mapping.items(): 

118 if quantity.endswith(suffix): 

119 # remove suffix, store multiplier, and break 

120 quantity = quantity[: -len(suffix)].strip() 

121 multiplier = mult 

122 break 

123 else: 

124 raise ValueError(f"Invalid suffix in {quantity_original}") 

125 else: 

126 # multiple suffixes 

127 raise ValueError(f"Multiple suffixes detected in {quantity_original}") 

128 

129 # fractions 

130 if "/" in quantity: 

131 try: 

132 assert quantity.count("/") == 1, "too many '/'" 

133 # split and strip 

134 num, den = quantity.split("/") 

135 num = num.strip() 

136 den = den.strip() 

137 num_sign: int = 1 

138 # negative numbers 

139 if num.startswith("-"): 

140 num_sign = -1 

141 num = num[1:] 

142 # assert that both are digits 

143 assert ( 

144 num.isdigit() and den.isdigit() 

145 ), "numerator and denominator must be digits" 

146 # return the fraction 

147 result = num_sign * ( 

148 int(num) / int(den) 

149 ) # this allows for fractions with suffixes, which is weird, but whatever 

150 except AssertionError as e: 

151 raise ValueError(f"Invalid fraction {quantity_original}: {e}") from e 

152 

153 # decimals 

154 else: 

155 try: 

156 result = int(quantity) 

157 except ValueError: 

158 try: 

159 result = float(quantity) 

160 except ValueError as e: 

161 raise ValueError( 

162 f"Invalid quantity {quantity_original} ({quantity})" 

163 ) from e 

164 

165 return result * multiplier