✨ Feature: add lock file for codegen (#231)

Author: yanyongyu

codegen/__init__.py

@@ -1,10 +1,9 @@
 from pathlib import Path
 import shutil
-import sys
 from typing import Any
 
-import httpx
 from jinja2 import Environment, PackageLoader
+import tomlkit
 
 from .config import Config
 from .log import logger as logger
@@ -21,10 +20,7 @@
 from .parser.schemas import UnionSchema
 from .source import get_source
 
-if sys.version_info >= (3, 11):
-    import tomllib
-else:
-    import tomli as tomllib
+LOCK_FILE_NAME = "versions.lock"
 
 env = Environment(
     loader=PackageLoader("codegen"),
@@ -46,7 +42,9 @@
 
 
 def load_config() -> Config:
-    pyproject = tomllib.loads(Path("./pyproject.toml").read_text(encoding="utf-8"))
+    pyproject = tomlkit.parse(
+        Path("./pyproject.toml").read_text(encoding="utf-8")
+    ).unwrap()
     config_dict: dict[str, Any] = pyproject.get("tool", {}).get("codegen", {})
 
     return Config.model_validate(config_dict)
@@ -252,6 +250,10 @@ def build_versions(dir: Path, versions: dict[str, str], latest_version: str):
     logger.info("Successfully generated versions!")
 
 
+def build_lock_file(file: Path, lock_data: tomlkit.TOMLDocument):
+    file.write_text(lock_data.as_string())
+
+
 def build():
     config = load_config()
     logger.info(f"Loaded config: {config!r}")
@@ -277,7 +279,8 @@ def build():
 
     for description in config.descriptions:
         logger.info(f"Start getting OpenAPI source for {description.identifier}...")
-        source = get_source(httpx.URL(description.source))
+        source = get_source(description.source)
+        description._actual_source = str(source.uri.copy_with(fragment=None))
         logger.info(f"Getting schema from {source.uri} succeeded!")
 
         logger.info(f"Start parsing OpenAPI spec for {description.identifier}...")
@@ -353,3 +356,4 @@ def build():
         versions[latest_version],
         latest_model_names,
     )
+    build_lock_file(config.output_dir / LOCK_FILE_NAME, config.to_lock())

codegen/config.py

@@ -2,6 +2,9 @@
 from typing import Any
 
 from pydantic import BaseModel, Field
+from tomlkit import aot, comment, document, inline_table, item, table
+from tomlkit.items import Table
+from tomlkit.toml_document import TOMLDocument
 
 
 class Override(BaseModel):
@@ -13,6 +16,29 @@ class Override(BaseModel):
 class VersionedOverride(Override):
     target_descriptions: list[str] = Field(default_factory=list)
 
+    def to_lock(self) -> Table:
+        tab = table()
+        if self.target_descriptions:
+            tab.append("target_descriptions", item(self.target_descriptions))
+        if self.class_overrides:
+            class_tab = table()
+            for key, value in self.class_overrides.items():
+                class_tab.append(key, value)
+            tab.append("class_overrides", class_tab)
+        if self.field_overrides:
+            field_tab = table()
+            for key, value in self.field_overrides.items():
+                field_tab.append(key, value)
+            tab.append("field_overrides", field_tab)
+        if self.schema_overrides:
+            schema_tab = table()
+            for key, value in self.schema_overrides.items():
+                value_table = inline_table()
+                value_table.update(value)
+                schema_tab.append(key, value_table)
+            tab.append("schema_overrides", schema_tab)
+        return tab
+
 
 class DescriptionConfig(BaseModel):
     version: str
@@ -24,11 +50,31 @@ class DescriptionConfig(BaseModel):
     is_latest: bool = False
     """If true, the description will be used as the default description."""
     source: str
+    """Source link to the description file."""
 
+    _actual_source: str | None = None
+    """The actual source link after downloading, if applicable."""
 
-class Config(BaseModel):
-    output_dir: Path
-    legacy_rest_models: Path
+    @property
+    def actual_source(self) -> str:
+        """Returns the actual source link after downloading, if applicable."""
+        return self._actual_source or self.source
+
+    def to_lock(self) -> Table:
+        tab = table()
+        tab.update(
+            {
+                "version": self.version,
+                "identifier": self.identifier,
+                "module": self.module,
+                "is_latest": self.is_latest,
+                "source": self.actual_source,
+            }
+        )
+        return tab
+
+
+class GenerationInfo(BaseModel):
     descriptions: list[DescriptionConfig]
     overrides: list[VersionedOverride] = Field(default_factory=list)
 
@@ -56,3 +102,24 @@ def get_override_config_for_version(self, version_id: str) -> Override:
                 for key, value in override.schema_overrides.items()
             },
         )
+
+    def to_lock(self) -> TOMLDocument:
+        doc = document()
+        doc.append(None, comment("DO NOT EDIT THIS FILE!"))
+        doc.append(None, comment("This file is automatically @generated by githubkit."))
+
+        descriptions_aot = aot()
+        for description in self.descriptions:
+            descriptions_aot.append(description.to_lock())
+        doc.append("descriptions", descriptions_aot)
+
+        overrides_aot = aot()
+        for override in self.overrides:
+            overrides_aot.append(override.to_lock())
+        doc.append("overrides", overrides_aot)
+        return doc
+
+
+class Config(GenerationInfo):
+    output_dir: Path
+    legacy_rest_models: Path

codegen/source.py

@@ -1,12 +1,20 @@
 from dataclasses import dataclass
 from functools import cache
-import json
-from pathlib import Path
+import os
 from typing import Any
 
 import httpx
 from jsonpointer import JsonPointer
 
+REPO_COMMIT = os.getenv(
+    "REPO_COMMIT",
+    "https://api.github.com/repos/github/rest-api-description/commits/main",
+)
+RAW_SOURCE_PREFIX = os.getenv(
+    "RAW_SOURCE_PREFIX",
+    "https://raw.github.com/github/rest-api-description/",
+)
+
 
 @dataclass(frozen=True)
 class Source:
@@ -37,19 +45,32 @@ def __truediv__(self, other: str | int) -> "Source":
 
 
 @cache
-def get_content(source: httpx.URL | Path) -> dict:
-    return (
-        json.loads(source.read_text(encoding="utf-8"))
-        if isinstance(source, Path)
-        else httpx.get(
-            source, headers={"User-Agent": "GitHubKit Codegen"}, follow_redirects=True
-        ).json()
+def get_content(source: str | httpx.URL) -> tuple[httpx.URL, dict]:
+    if isinstance(source, str):
+        sha_response = httpx.get(
+            REPO_COMMIT,
+            headers={
+                "User-Agent": "GitHubKit Codegen",
+                "Accept": "application/vnd.github.sha",
+            },
+        )
+        sha_response.raise_for_status()
+        sha = sha_response.text.strip()
+        source_link = httpx.URL(RAW_SOURCE_PREFIX).join(f"{sha}/{source.lstrip('/')}")
+    else:
+        source_link = source
+
+    response = httpx.get(
+        source_link, headers={"User-Agent": "GitHubKit Codegen"}, follow_redirects=True
     )
+    response.raise_for_status()
+    uri = response.url
+    content = response.json()
+    return uri, content
 
 
-def get_source(source: httpx.URL | Path, path: str | None = None) -> Source:
-    if isinstance(source, Path):
-        uri = httpx.URL(source.resolve().as_uri(), fragment=path)
-    else:
-        uri = source if path is None else source.copy_with(fragment=path)
-    return Source(uri=uri, root=get_content(source))
+def get_source(source: str | httpx.URL, path: str | None = None) -> Source:
+    uri, root = get_content(source)
+    if path is not None:
+        uri = uri.copy_with(fragment=path)
+    return Source(uri=uri, root=root)