1 from __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

22 def 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

49 def 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