Coverage for src/iso_freeze/iso_freeze.py: 63%

142 statements  

« prev     ^ index     » next       coverage.py v6.4.2, created at 2022-07-25 23:13 +0200

1"""Use pip install --report flag to separate pinned requirements for 

2different optional dependencies (e.g. 'dev' and 'doc' requirements).""" 

3 

4import argparse 

5import subprocess 

6import sys 

7import json 

8from dataclasses import dataclass 

9from typing import Any, Optional, Union 

10from pathlib import Path 

11 

12if sys.version_info >= (3, 11, 0): 

13 import tomllib 

14else: 

15 import tomli as tomllib # type: ignore 

16 

17 

18@dataclass 

19class PyPackage: 

20 """Class to capture relevant information about Python packages.""" 

21 

22 name: str 

23 version: str 

24 requested: bool = False 

25 hash: Optional[str] = None 

26 

27 

28def read_toml( 

29 toml_file: Path, 

30 optional_dependency: Optional[str] = None, 

31) -> list[str]: 

32 """Read TOML file and return list dependencies. 

33 

34 Includes requirements for optional dependency if any has been specified. 

35 

36 Keyword Arguments: 

37 toml_file -- Path to pyproject.toml file 

38 optional_dependency -- Optional dependency to include (default: None) 

39 

40 Returns: 

41 List of dependency names (list[str]) 

42 """ 

43 with open(toml_file, "rb") as f: 

44 metadata = tomllib.load(f) 

45 if not metadata.get("project"): 

46 sys.exit("TOML file does not contain a 'project' section.") 

47 dependencies: list[str] = metadata["project"].get("dependencies") 

48 if optional_dependency: 

49 if not metadata["project"].get("optional-dependencies"): 

50 sys.exit("No optional dependencies defined in TOML file.") 

51 optional_dependency_reqs: Optional[list[str]] = ( 

52 metadata["project"].get("optional-dependencies").get(optional_dependency) 

53 ) 

54 if optional_dependency_reqs: 

55 dependencies.extend(optional_dependency_reqs) 

56 else: 

57 sys.exit( 

58 f"No optional dependency '{optional_dependency}' found in TOML file." 

59 ) 

60 return dependencies 

61 

62 

63def build_pip_report_command( 

64 python_exec: Path, 

65 toml_dependencies: Optional[list[str]], 

66 requirements_in: Optional[Path], 

67 pip_args: Optional[list[str]], 

68) -> list[Union[str, Path]]: 

69 """Build pip command to to generate report. 

70 

71 Arguments: 

72 python_exec -- Path to Python interpreter to use (Path) 

73 toml_dependencies -- TOML dependencies to install (Optional[list[str]]) 

74 requirements_in -- Path to requirements file (Optional[Path]) 

75 pip_args -- Arguments to be passed to pip install (Optional[list[str]]) 

76 

77 Returns: 

78 Pip command to pass to run_pip_report (list[Union[str, Path]]) 

79 """ 

80 pip_report_command: list[Union[str, Path]] = [ 

81 "env", 

82 "PIP_REQUIRE_VIRTUALENV=false", 

83 python_exec, 

84 "-m", 

85 "pip", 

86 "install", 

87 ] 

88 # If pip_args have been provided, inject them after the 'install' keyword 

89 if pip_args: 

90 pip_report_command.extend(pip_args) 

91 # Add necessary flags for calling pip install report 

92 pip_report_command.extend( 

93 ["-q", "--dry-run", "--ignore-installed", "--report", "-"] 

94 ) 

95 # Finally, either append dependencies from TOML file or '-r requirements-file' 

96 if toml_dependencies: 

97 pip_report_command.extend([dependency for dependency in toml_dependencies]) 

98 elif requirements_in: 

99 pip_report_command.extend(["-r", requirements_in]) 

100 return pip_report_command 

101 

102 

103def run_pip(command: list[Union[str, Path]], check_output: bool) -> Any: 

104 """Run specified pip command with subprocess and return results, if any. 

105 

106 Arguments: 

107 command -- pip command to execute (list[Union[str, Path]]) 

108 

109 Keyword Arguments: 

110 check_output -- Whether to call subprocess.check_output (default: {False}) 

111 

112 Returns: 

113 Output of pip command, if any (Optional[Any]) 

114 """ 

115 try: 

116 if check_output: 

117 return subprocess.check_output(command, encoding="utf-8") 

118 else: 

119 subprocess.run(command) 

120 return None 

121 except subprocess.CalledProcessError as error: 

122 error.output 

123 sys.exit() 

124 

125 

126def get_dependencies( 

127 pip_report_command: list[Union[Path, str]], 

128) -> Optional[list[PyPackage]]: 

129 """Capture pip install --report to generate pinned requirements. 

130 

131 Arguments: 

132 pip_report_command -- Command for subprocess (list[Union[Path, str]]) 

133 

134 Returns: 

135 List of PyPackage objects containing infos to pin requirements (list[PyPackage]) 

136 """ 

137 pip_report: dict[str, Any] = json.loads( 

138 run_pip(command=pip_report_command, check_output=True) 

139 ) 

140 if pip_report.get("install"): 

141 dependencies: list[PyPackage] = [] 

142 for package in pip_report["install"]: 

143 dependencies.append( 

144 PyPackage( 

145 name=package["metadata"]["name"], 

146 version=package["metadata"]["version"], 

147 requested=package["requested"], 

148 # pip report provides hashes in the form 'sha256=<hash>', but pip 

149 # install requires 'sha256:<hash>', so we replace '=' with ':' 

150 hash=package["download_info"]["archive_info"]["hash"].replace( 

151 "=", ":" 

152 ), 

153 ) 

154 ) 

155 return dependencies 

156 else: 

157 return None 

158 

159 

160def get_installed_packages(python_exec: Path) -> list[PyPackage]: 

161 """Run pip list --format json and return packages. 

162 

163 Returns: 

164 List of packages from current environment (list[PyPackage]) 

165 """ 

166 pip_list_output: list[dict[str, str]] = json.loads( 

167 run_pip( 

168 command=[ 

169 python_exec, 

170 "-m", 

171 "pip", 

172 "list", 

173 "--format", 

174 "json", 

175 "--exclude-editable", 

176 ], 

177 check_output=True, 

178 ) 

179 ) 

180 installed_packages: list[PyPackage] = [] 

181 for package in pip_list_output: 

182 installed_packages.append( 

183 PyPackage(name=package["name"], version=package["version"]) 

184 ) 

185 return installed_packages 

186 

187 

188def remove_additional_packages( 

189 installed_packages: list[PyPackage], to_install: list[PyPackage], python_exec: Path 

190) -> None: 

191 """Remove packages installed in environment but not included in pip install --report. 

192 

193 Arguments: 

194 installed_packages -- List of packages installed in current environment 

195 (list[PyPackage]) 

196 to_install -- List of packages taken from pip install --report 

197 python_exec -- Path to Python executable (Path) 

198 """ 

199 # Create two lists with packages names only for easy comparison 

200 installed_names_only: list[str] = [package.name for package in installed_packages] 

201 to_install_names_only: list[str] = [package.name for package in to_install] 

202 to_delete: list[str] = [ 

203 package 

204 for package in installed_names_only 

205 if package not in to_install_names_only 

206 # Don't remove default packages 

207 if package not in ["pip", "setuptools"] 

208 ] 

209 if to_delete: 

210 run_pip( 

211 command=[python_exec, "-m", "pip", "uninstall", "-y", *to_delete], 

212 check_output=False, 

213 ) 

214 

215 

216def install_pip_report_output(to_install: list[PyPackage], python_exec: Path) -> None: 

217 """Install packages with pinned versions from pip install --report output. 

218 

219 Arguments: 

220 to_install -- List of packages taken from pip install --report 

221 python_exec -- Path to Python executable (Path) 

222 """ 

223 # Create list in the format ["package1==versionX", "package2==versionY"] 

224 # from `pip install --report` and pass that to `pip install` 

225 pinned_versions: list[str] = [ 

226 f"{package.name}=={package.version}" for package in to_install 

227 ] 

228 run_pip( 

229 command=[python_exec, "-m", "pip", "install", "--upgrade", *pinned_versions], 

230 check_output=False, 

231 ) 

232 

233 

234def build_reqirements_file_contents( 

235 dependencies: list[PyPackage], hashes: bool 

236) -> None: 

237 """Build lists to be written to a requirements file. 

238 

239 Display top level dependencies on top, similar to pip freeze -r requirements_file. 

240 

241 Arguments: 

242 dependencies -- Dependencies listed in pip install --report (list[dict[str]]) 

243 hashes -- Whether to include hashes (bool) 

244 

245 Returns: 

246 Contents of requirements file (list[str]) 

247 """ 

248 # For easier formatting we create separate lists for top level requirements 

249 # and their dependencies 

250 top_level_requirements: list[str] = [] 

251 dependency_requirements: list[str] = [] 

252 for package in dependencies: 

253 pinned_format: str = f"{package.name}=={package.version}" 

254 if hashes: 

255 pinned_format += f" \\\n --hash={package.hash}" 

256 # If requested == True, the package is a top level requirement 

257 if package.requested: 

258 top_level_requirements.append(pinned_format) 

259 else: 

260 dependency_requirements.append(pinned_format) 

261 # Sort pinned packages alphabetically before writing to file 

262 # (case-insensitively thanks to key=str.lower) 

263 top_level_requirements.sort(key=str.lower) 

264 # Combine lists and add comments 

265 requirements_file_content: list[str] = [ 

266 "# Top level requirements", 

267 *top_level_requirements, 

268 ] 

269 if dependency_requirements: 

270 dependency_requirements.sort(key=str.lower) 

271 requirements_file_content.extend( 

272 ["# Dependencies of top level requirements", *dependency_requirements] 

273 ) 

274 return requirements_file_content 

275 

276 

277def write_requirements_file(output_file: Path, file_contents: list[str]) -> None: 

278 """Write requirements file. 

279 

280 Arguments: 

281 output_file -- Path to and name of requirements.txt file (Path) 

282 file_contents -- Contents to to written to a file (list[str]) 

283 """ 

284 with output_file.open(mode="w", encoding="utf=8") as f: 

285 f.writelines(f"{package}\n" for package in file_contents) 

286 

287 

288def determine_default_file() -> Optional[Path]: 

289 """Determine default input file if none has been specified. 

290 

291 Returns: 

292 Path to default file (Optional[Path]) 

293 """ 

294 if Path("requirements.in").exists(): 

295 default: Optional[Path] = Path("requirements.in") 

296 elif Path("pyproject.toml").exists(): 

297 default = Path("pyproject.toml") 

298 else: 

299 default = None 

300 return default 

301 

302 

303def parse_args() -> argparse.Namespace: 

304 """Parse arguments.""" 

305 argparser = argparse.ArgumentParser( 

306 description="Use pip install --report to cleanly separate pinned requirements " 

307 "for different optional dependencies (e.g. 'dev' and 'doc' requirements)." 

308 ) 

309 argparser.add_argument( 

310 "file", 

311 type=Path, 

312 nargs="?", 

313 default=determine_default_file(), 

314 help="Path to input file. Can be pyproject.toml or requirements file. " 

315 "Defaults to 'requirements.in' or 'pyproject.toml' in current directory.", 

316 ) 

317 argparser.add_argument( 

318 "--dependency", 

319 "-d", 

320 type=str, 

321 help="Name of the optional dependency defined in pyproject.toml to include.", 

322 ) 

323 argparser.add_argument( 

324 "--output", 

325 "-o", 

326 type=Path, 

327 default=Path("requirements.txt"), 

328 help="Name of the output file. Defaults to 'requirements.txt' if unspecified.", 

329 ) 

330 argparser.add_argument( 

331 "--python", 

332 "-p", 

333 type=Path, 

334 default=Path("python3"), 

335 help="Specify path to Python interpreter to use. Defaults to 'python3'.", 

336 ) 

337 argparser.add_argument( 

338 "--pip-args", 

339 type=str, 

340 help="List of arguments to be passed to pip install. Call as: " 

341 'pip-args "--arg1 value --arg2 value".', 

342 ) 

343 argparser.add_argument( 

344 "--sync", 

345 "-s", 

346 action="store_true", 

347 help="Sync current environment with dependencies listed in file (removes " 

348 "packages that are not dependencies in file, adds those that are missing)", 

349 ) 

350 argparser.add_argument( 

351 "--hashes", action="store_true", help="Add hashes to output file." 

352 ) 

353 args = argparser.parse_args() 

354 if not args.file: 

355 sys.exit( 

356 "No requirements.in or pyproject.toml file found in current directory. " 

357 "Please specify input file." 

358 ) 

359 if args.file.suffix != ".toml" and args.dependency: 

360 sys.exit( 

361 "You can only specify an optional dependency if your input file is " 

362 "pyproject.toml." 

363 ) 

364 if not args.file.is_file(): 

365 sys.exit(f"Not a file: {args.file}") 

366 # If pip-args have been provided, split them into list 

367 if args.pip_args: 

368 args.pip_args = args.pip_args.split(" ") 

369 return args 

370 

371 

372def validate_pip_version(python_exec: Path) -> bool: 

373 """Check if pip version is >= 22.2. 

374 

375 Returns: 

376 True/False (bool) 

377 """ 

378 pip_version_call: str = run_pip( 

379 command=[python_exec, "-m", "pip", "--version"], check_output=True 

380 ) 

381 # Output of pip --version looks like this: 

382 # pip 22.2 from <path to pip> (<python version>) 

383 # To get version number, split this message on whitespace and pick list item 1. 

384 # To check against minimum version, turn the version number into a list of ints 

385 # (e.g. '[22, 2]' or '[21, 1, 2]') 

386 pip_version: list[int] = [ 

387 int(number) for number in pip_version_call.split()[1].split(".") 

388 ] 

389 if pip_version >= [22, 2]: 

390 return True 

391 return False 

392 

393 

394def main() -> None: 

395 arguments: argparse.Namespace = parse_args() 

396 if not validate_pip_version(arguments.python): 

397 sys.exit("pip >= 22.2 required. Please update pip and try again.") 

398 if arguments.file.suffix == ".toml": 

399 toml_dependencies: Optional[list[str]] = read_toml( 

400 toml_file=arguments.file, optional_dependency=arguments.dependency 

401 ) 

402 else: 

403 toml_dependencies = None 

404 pip_report_command: list[Union[str, Path]] = build_pip_report_command( 

405 python_exec=arguments.python, 

406 toml_dependencies=toml_dependencies, 

407 requirements_in=arguments.file, 

408 pip_args=arguments.pip_args, 

409 ) 

410 dependencies: Optional[list[PyPackage]] = get_dependencies( 

411 pip_report_command=pip_report_command 

412 ) 

413 if dependencies: 

414 if arguments.sync: 

415 remove_additional_packages( 

416 installed_packages=get_installed_packages(python_exec=arguments.python), 

417 to_install=dependencies, 

418 python_exec=arguments.python, 

419 ) 

420 install_pip_report_output( 

421 to_install=dependencies, 

422 python_exec=arguments.python, 

423 ) 

424 else: 

425 output_file_contents: list[str] = build_reqirements_file_contents( 

426 dependencies=dependencies, hashes=arguments.hashes 

427 ) 

428 write_requirements_file( 

429 output_file=arguments.output, file_contents=output_file_contents 

430 ) 

431 print(f"Pinned specified requirements in {arguments.output}") 

432 else: 

433 sys.exit("There are no dependencies to pin") 

434 

435 

436if __name__ == "__main__": 

437 main()