pylock.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188
  1. from __future__ import annotations
  2. import dataclasses
  3. import re
  4. from collections.abc import Iterable
  5. from dataclasses import dataclass
  6. from pathlib import Path
  7. from typing import TYPE_CHECKING, Any
  8. from pip._vendor import tomli_w
  9. from pip._internal.models.direct_url import ArchiveInfo, DirInfo, VcsInfo
  10. from pip._internal.models.link import Link
  11. from pip._internal.req.req_install import InstallRequirement
  12. from pip._internal.utils.urls import url_to_path
  13. if TYPE_CHECKING:
  14. from typing_extensions import Self
  15. PYLOCK_FILE_NAME_RE = re.compile(r"^pylock\.([^.]+)\.toml$")
  16. def is_valid_pylock_file_name(path: Path) -> bool:
  17. return path.name == "pylock.toml" or bool(re.match(PYLOCK_FILE_NAME_RE, path.name))
  18. def _toml_dict_factory(data: list[tuple[str, Any]]) -> dict[str, Any]:
  19. return {key.replace("_", "-"): value for key, value in data if value is not None}
  20. @dataclass
  21. class PackageVcs:
  22. type: str
  23. url: str | None
  24. # (not supported) path: Optional[str]
  25. requested_revision: str | None
  26. commit_id: str
  27. subdirectory: str | None
  28. @dataclass
  29. class PackageDirectory:
  30. path: str
  31. editable: bool | None
  32. subdirectory: str | None
  33. @dataclass
  34. class PackageArchive:
  35. url: str | None
  36. # (not supported) path: Optional[str]
  37. # (not supported) size: Optional[int]
  38. # (not supported) upload_time: Optional[datetime]
  39. hashes: dict[str, str]
  40. subdirectory: str | None
  41. @dataclass
  42. class PackageSdist:
  43. name: str
  44. # (not supported) upload_time: Optional[datetime]
  45. url: str | None
  46. # (not supported) path: Optional[str]
  47. # (not supported) size: Optional[int]
  48. hashes: dict[str, str]
  49. @dataclass
  50. class PackageWheel:
  51. name: str
  52. # (not supported) upload_time: Optional[datetime]
  53. url: str | None
  54. # (not supported) path: Optional[str]
  55. # (not supported) size: Optional[int]
  56. hashes: dict[str, str]
  57. @dataclass
  58. class Package:
  59. name: str
  60. version: str | None = None
  61. # (not supported) marker: Optional[str]
  62. # (not supported) requires_python: Optional[str]
  63. # (not supported) dependencies
  64. vcs: PackageVcs | None = None
  65. directory: PackageDirectory | None = None
  66. archive: PackageArchive | None = None
  67. # (not supported) index: Optional[str]
  68. sdist: PackageSdist | None = None
  69. wheels: list[PackageWheel] | None = None
  70. # (not supported) attestation_identities: Optional[List[Dict[str, Any]]]
  71. # (not supported) tool: Optional[Dict[str, Any]]
  72. @classmethod
  73. def from_install_requirement(cls, ireq: InstallRequirement, base_dir: Path) -> Self:
  74. base_dir = base_dir.resolve()
  75. dist = ireq.get_dist()
  76. download_info = ireq.download_info
  77. assert download_info
  78. package = cls(name=dist.canonical_name)
  79. if ireq.is_direct:
  80. if isinstance(download_info.info, VcsInfo):
  81. package.vcs = PackageVcs(
  82. type=download_info.info.vcs,
  83. url=download_info.url,
  84. requested_revision=download_info.info.requested_revision,
  85. commit_id=download_info.info.commit_id,
  86. subdirectory=download_info.subdirectory,
  87. )
  88. elif isinstance(download_info.info, DirInfo):
  89. package.directory = PackageDirectory(
  90. path=(
  91. Path(url_to_path(download_info.url))
  92. .resolve()
  93. .relative_to(base_dir)
  94. .as_posix()
  95. ),
  96. editable=(
  97. download_info.info.editable
  98. if download_info.info.editable
  99. else None
  100. ),
  101. subdirectory=download_info.subdirectory,
  102. )
  103. elif isinstance(download_info.info, ArchiveInfo):
  104. if not download_info.info.hashes:
  105. raise NotImplementedError()
  106. package.archive = PackageArchive(
  107. url=download_info.url,
  108. hashes=download_info.info.hashes,
  109. subdirectory=download_info.subdirectory,
  110. )
  111. else:
  112. # should never happen
  113. raise NotImplementedError()
  114. else:
  115. package.version = str(dist.version)
  116. if isinstance(download_info.info, ArchiveInfo):
  117. if not download_info.info.hashes:
  118. raise NotImplementedError()
  119. link = Link(download_info.url)
  120. if link.is_wheel:
  121. package.wheels = [
  122. PackageWheel(
  123. name=link.filename,
  124. url=download_info.url,
  125. hashes=download_info.info.hashes,
  126. )
  127. ]
  128. else:
  129. package.sdist = PackageSdist(
  130. name=link.filename,
  131. url=download_info.url,
  132. hashes=download_info.info.hashes,
  133. )
  134. else:
  135. # should never happen
  136. raise NotImplementedError()
  137. return package
  138. @dataclass
  139. class Pylock:
  140. lock_version: str = "1.0"
  141. # (not supported) environments: Optional[List[str]]
  142. # (not supported) requires_python: Optional[str]
  143. # (not supported) extras: List[str] = []
  144. # (not supported) dependency_groups: List[str] = []
  145. created_by: str = "pip"
  146. packages: list[Package] = dataclasses.field(default_factory=list)
  147. # (not supported) tool: Optional[Dict[str, Any]]
  148. def as_toml(self) -> str:
  149. return tomli_w.dumps(dataclasses.asdict(self, dict_factory=_toml_dict_factory))
  150. @classmethod
  151. def from_install_requirements(
  152. cls, install_requirements: Iterable[InstallRequirement], base_dir: Path
  153. ) -> Self:
  154. return cls(
  155. packages=sorted(
  156. (
  157. Package.from_install_requirement(ireq, base_dir)
  158. for ireq in install_requirements
  159. ),
  160. key=lambda p: p.name,
  161. )
  162. )