Coverage for src/pythia/utils/ext.py: 78%

1"""Python extensions and decorators.""" 

2 

3from __future__ import annotations 

4 

5import importlib.util 

6import platform 

7import sys 

8from importlib import import_module 

9from pathlib import Path 

10from types import ModuleType 

11from typing import Any 

12from typing import Iterable 

13from typing import Optional 

14from typing import Tuple 

15from typing import TypeVar 

16 

17RT = TypeVar("RT", Path, None) # return type 

18T = TypeVar("T") 

19 

20 

21def not_empty(path: Path) -> Path: 

22 """Raise if receives an empty value. 

23 

24 Args: 

25 path: The file to check for emptiness. 

26 

27 Returns: 

28 The decorated function. 

29 

30 Raises: 

31 EOFError: The received file is empty. 

32 

33 """ 

34 if path is None: 34 ↛ 35line 34 didn't jump to line 35, because the condition on line 34 was never true

35 return path 

36 if not path.stat().st_size: 36 ↛ 37line 36 didn't jump to line 37, because the condition on line 36 was never true

37 raise EOFError(f"File {path} is empty") 

38 return path 

39 

40 

41def not_none(value: Optional[Any]): 

42 """Raise if receives `None`. 

43 

44 Args: 

45 value: The value which should not be none. 

46 

47 Returns: 

48 The received value. 

49 

50 Raises: 

51 ValueError: The received value was `None`. 

52 

53 """ 

54 

55 if value is None: 55 ↛ 57line 55 didn't jump to line 57, because the condition on line 55 was never true

56 

57 raise ValueError("Received disallowed `None`") 

58 return value 

59 

60 

61def get_arch() -> str: 

62 """Return system arch. 

63 

64 Returns: 

65 platform, like `uname machine`. 

66 

67 """ 

68 return platform.uname()[4] 

69 

70 

71def import_from_path(name: str, path: str | Path) -> ModuleType: 

72 """Import a module from a filepath. 

73 

74 Args: 

75 name: the name to use for the module when importing. 

76 path: path to the python file to import as a module. 

77 

78 Returns: 

79 The imported module. 

80 

81 Raises: 

82 ImportError: unable to get spec from location 

83 

84 See Also: 

85 `importlib.util.spec_from_file_location` 

86 

87 """ 

88 spec = importlib.util.spec_from_file_location(name, str(path)) 

89 if not spec or not spec.loader: 89 ↛ 90line 89 didn't jump to line 90, because the condition on line 89 was never true

90 raise ImportError(f"Failed to import from {path=}") 

91 

92 module = importlib.util.module_from_spec(spec) 

93 sys.modules[name] = module 

94 spec.loader.exec_module(module) 

95 return module 

96 

97 

98def import_from_str(module: str, *, name=None, suffix: str = "") -> ModuleType: 

99 """Import python file as a module or from its path. 

100 

101 Args: 

102 module: name of the python file to import. 

103 name: ``__name__`` to use when import thing the module. If not 

104 set, defaults to the file's stem. 

105 suffix: file suffix. If set, load `name` as a file. Otherwise, 

106 

107 Returns: 

108 The improted module. 

109 

110 """ 

111 if suffix: 111 ↛ 114line 111 didn't jump to line 114, because the condition on line 111 was never false

112 path = Path(module).with_suffix(suffix) 

113 return import_from_path(name or path.stem, path) 

114 return import_module(module) 

115 

116 

117def grouped(iterable: Iterable[T], size=2) -> Iterable[Tuple[T, ...]]: 

118 """Iterate by groups. 

119 

120 Args: 

121 iterable: container for the data to group by. 

122 size: size of the groups 

123 

124 Returns: 

125 zip containing tuples of values. 

126 

127 Example: 

128 >>> [*grouped(range(10),3)] 

129 [(0, 1, 2), (3, 4, 5), (6, 7, 8)] 

130 

131 """ 

132 return zip(*[iter(iterable)] * size) 

133 

134 

135def remove_suffix(input_string: str, suffix: str) -> str: 

136 """Remove trailing substring. 

137 

138 Args: 

139 input_string: Input string to look for the suffix. 

140 suffix: The substring to match at the end of the input string. 

141 

142 Returns: 

143 If the string ends with the suffix string and that suffix 

144 is not empty, return string[:-len(suffix)]. Otherwise, return a 

145 copy of the original string 

146 

147 Backport of stdlib@3.9 

148 

149 """ 

150 if suffix and input_string.endswith(suffix): 150 ↛ 152line 150 didn't jump to line 152, because the condition on line 150 was never false

151 return input_string[: -len(suffix)] 

152 return input_string 

153 

154 

155def remove_prefix(input_string: str, prefix: str) -> str: 

156 """Remove leading substring. 

157 

158 Args: 

159 input_string: Input string to look for the suffix. 

160 prefix: The substring to match at the start of the input string. 

161 

162 Returns: 

163 If the string ends with the suffix string and that suffix 

164 is not empty, return string[:-len(suffix)]. Otherwise, return a 

165 copy of the original string 

166 

167 Backport of stdlib@3.9 

168 

169 """ 

170 if prefix and input_string.startswith(prefix): 170 ↛ 172line 170 didn't jump to line 172, because the condition on line 170 was never false

171 return input_string[len(prefix) :] 

172 return input_string 

173 

174 

175IS_JETSON = get_arch() == "aarch64"