Coverage for src/lektor_ng/packages.py: 73%

1from __future__ import annotations 

2 

3import hashlib 

4import os 

5import shutil 

6import site 

7import subprocess 

8import sys 

9import sysconfig 

10from collections.abc import Iterable, Iterator, Sized 

11from pathlib import Path 

12from typing import TYPE_CHECKING 

13from venv import EnvBuilder 

14 

15import click 

16import requests 

17 

18if TYPE_CHECKING: 

19 from _typeshed import StrPath 

20 from lektor.environment import Environment # circ dependency 

21else: 

22 StrPath = object 

23 

24 

25def _get_package_version_from_project(cfg, name): 

26 choices = (name.lower(), "lektor-" + name.lower()) 

27 for pkg, version in cfg.section_as_dict("packages").items(): 

28 if pkg.lower() in choices: 

29 return {"name": pkg, "version": version} 

30 return None 

31 

32 

33def add_package_to_project(project, req): 

34 """Given a package requirement this returns the information about this 

35 plugin. 

36 """ 

37 if "@" in req: 

38 name, version = req.split("@", 1) 

39 version_hint = version 

40 else: 

41 name = req 

42 version = None 

43 version_hint = "latest release" 

44 

45 cfg = project.open_config() 

46 info = _get_package_version_from_project(cfg, name) 

47 if info is not None: 

48 raise RuntimeError("The package was already added to the project.") 

49 

50 for choice in name, "lektor-" + name: 

51 rv = requests.get(f"https://pypi.python.org/pypi/{choice}/json", timeout=10) 

52 if rv.status_code != 200: 

53 continue 

54 

55 data = rv.json() 

56 canonical_name = data["info"]["name"] 

57 if version is None: 

58 version = data["info"]["version"] 

59 version_info = data["releases"].get(version) 

60 if version_info is None: 

61 raise RuntimeError(f"Latest requested version ({version_hint}) could not be found") 

62 

63 cfg[f"packages.{canonical_name}"] = version 

64 cfg.save() 

65 return {"name": canonical_name, "version": version} 

66 

67 raise RuntimeError("The package could not be found on PyPI") 

68 

69 

70def remove_package_from_project(project, name): 

71 cfg = project.open_config() 

72 choices = (name.lower(), "lektor-" + name.lower()) 

73 for pkg, version in cfg.section_as_dict("packages").items(): 

74 if pkg.lower() in choices: 

75 del cfg[f"packages.{pkg}"] 

76 cfg.save() 

77 return {"name": pkg, "version": version} 

78 return None 

79 

80 

81if os.name == "nt": 

82 _default_venv_symlinks = False 

83else: 

84 _default_venv_symlinks = True 

85 

86 

87class VirtualEnv: 

88 """A helper for manipulating our private package cache virtual environment. 

89 

90 Parameters: 

91 

92 path — The path to the virtual environment to manage. This can be an existing 

93 environment or not. 

94 

95 """ 

96 

97 def __init__(self, path: StrPath): 

98 self.path = Path(path) 

99 

100 def create( 

101 self, 

102 with_pip: bool = True, 

103 upgrade_deps: bool = True, 

104 symlinks: bool = _default_venv_symlinks, 

105 ) -> None: 

106 """(Re-)Create a new virtual environment. 

107 

108 This will remove any existing virtual environment and create a new one. 

109 

110 The parameters ``with_pip`` and ``upgrade_deps`` should probably be left at 

111 their default values in normal usage. They are provided here mostly for use in 

112 tests. They work as described for ``venv.EnvBuilder`` from the standard library. 

113 (Though ``upgrade_deps`` is only supported by EnvBuilder`` in py39+, here we 

114 emulate it's behavior if running under older pythons.) 

115 

116 """ 

117 # Right now, by default, we always install and upgrade pip to 

118 # the latest available version. 

119 # 

120 # We could optimize by not installing (and not upgrading) pip if 

121 # the system pip is sufficient to our needs. 

122 # 

123 # Note that, e.g., pip>=21.3 is required to support PEP660 editable 

124 # installs. 

125 env_builder = EnvBuilder( 

126 clear=True, 

127 with_pip=with_pip, 

128 upgrade_deps=upgrade_deps, 

129 symlinks=symlinks, 

130 ) 

131 env_builder.create(self.path) 

132 

133 def addsitedir(self, sitedir: str) -> None: 

134 """Add an additional sitedir to sys.path for virtual environment. 

135 

136 Packages installed in ``sitedir`` will be made available to any invocations 

137 of python running within the virtual environment. 

138 """ 

139 with Path(self.site_packages, "_lektor.pth").open("a", encoding="utf-8") as fp: 

140 fp.write(f"import site; site.addsitedir({sitedir!r})\n") 

141 

142 def run_pip_install(self, *args: str) -> None: 

143 """Run `pip install` in the virtual environment. 

144 

145 ``Args`` are appended to the command line (following ``pip install``). They 

146 should specify how and which packages to install. 

147 

148 """ 

149 try: 

150 subprocess.run((self.executable, "-m", "pip", "install", *args), check=True) 

151 except subprocess.CalledProcessError as exc: 

152 raise RuntimeError("Failed to install dependency package.") from exc 

153 

154 @property 

155 def site_packages(self) -> str: 

156 """The path to the virtual environments ``site-packages`` directory.""" 

157 return self._get_path("purelib") 

158 

159 @property 

160 def executable(self) -> str: 

161 """The path to the python interpreter for the virtual environment.""" 

162 script_path = Path(self._get_path("scripts")) 

163 executable_name = Path(sys.executable).name 

164 return os.fspath(script_path / executable_name) 

165 

166 def _get_path(self, name: str) -> str: 

167 vars = {"base": os.fspath(self.path)} 

168 return sysconfig.get_path(name, vars=vars) 

169 

170 

171class Requirements(Iterable[str], Sized): 

172 """Manage package requirements.""" 

173 

174 requirements: set[str] 

175 

176 def __init__(self) -> None: 

177 self.requirements = set() 

178 

179 def __len__(self) -> int: 

180 return len(self.requirements) 

181 

182 def __iter__(self) -> Iterator[str]: 

183 """The requirements. 

184 

185 These requirements are in the form of arguments that can be passed to ``pip 

186 install``. 

187 """ 

188 return iter(self.requirements) 

189 

190 def add_requirement(self, package: str, version: str | None = None) -> None: 

191 """Add a (remote) distribution to the requirements.""" 

192 self.requirements.add(f"{package}=={version}" if version else f"{package}") 

193 

194 def add_local_requirement(self, path: StrPath) -> None: 

195 """Add a local distribution source directory to the requirements. 

196 

197 The distribution source at ``path`` (which should be a legacy `setup.py` or 

198 modern PEP660-compatible project) will be installed in editable mode. 

199 

200 """ 

201 srcdir = os.fspath(Path(path).resolve()) 

202 self.requirements.add(f"--editable={srcdir}") 

203 

204 _DIST_FILES = ("setup.py", "pyproject.toml") 

205 

206 def add_local_requirements_from(self, packages_path: StrPath) -> None: 

207 """Add sub-directories of path that look like local distribution sources. 

208 

209 Any direct sub-directories of ``packages_path`` which appear to be distribution 

210 source code will be added to the requirements in local (editable) mode. 

211 

212 """ 

213 try: 

214 for path in Path(packages_path).iterdir(): 

215 if any(path.joinpath(fn).is_file() for fn in self._DIST_FILES): 

216 self.add_local_requirement(path) 

217 except OSError: 

218 pass 

219 

220 def hash(self) -> str: 

221 """Compute a hash of the requirement set.""" 

222 hash = hashlib.sha1() 

223 for requirement in sorted(self.requirements): 

224 hash.update(requirement.encode("utf-8")) 

225 hash.update(b"\0") 

226 return hash.hexdigest() 

227 

228 

229def update_cache( 

230 venv_path: Path, 

231 remote_packages: dict[str, str], 

232 local_package_path: Path, 

233) -> None: 

234 """Ensure the package cache at venv_path is up-to-date. 

235 

236 ``Remote_packages`` is a dictionary (mapping package names to required versions) 

237 that specifies remote packages (to be installed from PyPI). 

238 

239 ``Local_package_page`` is a path to a directory whose sub-directories may contain 

240 local plugin source. Any such source directories will be installed in "editable" 

241 mode. 

242 

243 """ 

244 requirements = Requirements() 

245 for package, version in remote_packages.items(): 

246 requirements.add_requirement(package, version) 

247 requirements.add_local_requirements_from(local_package_path) 

248 

249 if len(requirements) == 0: 

250 shutil.rmtree(venv_path, ignore_errors=True) 

251 else: 

252 hash_file = venv_path / "lektor-requirements-hash.txt" 

253 try: 

254 is_stale = hash_file.read_text().strip() != requirements.hash() 

255 except FileNotFoundError: 

256 is_stale = True 

257 

258 if is_stale: 

259 venv = VirtualEnv(venv_path) 

260 venv.create() 

261 # Add our site-packages to venv's sys.path 

262 our_site_packages = sysconfig.get_path("purelib") 

263 venv.addsitedir(our_site_packages) 

264 

265 venv.run_pip_install(*requirements) 

266 hash_file.write_text(f"{requirements.hash()}\n", encoding="ascii") 

267 

268 

269def load_packages(env: Environment, reinstall: bool = False) -> None: 

270 """Import all of our managed plugins into our ``sys.path`` 

271 

272 This first ensures that our private package cache is up-to-date, then 

273 adds it to ``sys.path``. 

274 

275 After ``load_packages`` is called, the entry points defined in by 

276 plugins that we manage will be available. 

277 """ 

278 if reinstall: 

279 click.echo("Force package cache refresh.") 

280 wipe_package_cache(env) 

281 

282 config = env.load_config() 

283 venv_path = env.project.get_package_cache_path() 

284 update_cache(venv_path, config["PACKAGES"], Path(env.root_path, "packages")) 

285 site.addsitedir(VirtualEnv(venv_path).site_packages) 

286 

287 

288def wipe_package_cache(env: Environment) -> None: 

289 """Remove the entire package cache.""" 

290 project = env.project 

291 # Remove the legacy flat package cache, too 

292 for cache_type in project.PackageCacheType: 

293 shutil.rmtree(project.get_package_cache_path(cache_type), ignore_errors=True)