|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +import argparse |
| 4 | +import json |
| 5 | +import logging |
| 6 | +import os |
| 7 | +import subprocess |
| 8 | +from pathlib import Path |
| 9 | +from venv import logger |
| 10 | + |
| 11 | +import httpx |
| 12 | +from httpx import HTTPStatusError, RequestError |
| 13 | +from rich.logging import RichHandler |
| 14 | +from rich.traceback import install |
| 15 | + |
| 16 | +# Replace the basic logging config with Rich handler |
| 17 | +logging.basicConfig( |
| 18 | + level=logging.INFO, |
| 19 | + format="%(message)s", |
| 20 | + handlers=[RichHandler(rich_tracebacks=True, markup=True)], |
| 21 | +) |
| 22 | + |
| 23 | +logging.getLogger("httpx").setLevel(logging.INFO) |
| 24 | + |
| 25 | +install(show_locals=False) |
| 26 | + |
| 27 | +logger = logging.getLogger(__name__) |
| 28 | + |
| 29 | + |
| 30 | +class OpenAPIDiscovery: |
| 31 | + def __init__( |
| 32 | + self, |
| 33 | + input_dir: str, |
| 34 | + output_dir: str, |
| 35 | + version: str, |
| 36 | + verbose: bool, |
| 37 | + use_cache: bool, |
| 38 | + ): |
| 39 | + self.base_url = "https://api.github.com/repos/eda-labs/openapi" |
| 40 | + self.raw_base = f"https://raw.githubusercontent.com/eda-labs/openapi/{version}" |
| 41 | + self.gh_token = os.environ.get("GH_AUTH_TOKEN") |
| 42 | + if self.gh_token == "": |
| 43 | + logger.error("GH_AUTH_TOKEN environment variable is not set.") |
| 44 | + exit(1) |
| 45 | + |
| 46 | + self.headers = { |
| 47 | + "Accept": "application/vnd.github.v3+json", |
| 48 | + "Authorization": f"Bearer {self.gh_token}", |
| 49 | + } |
| 50 | + self.cache_file = Path("cached_specs.json") |
| 51 | + self.input_dir = Path(input_dir) |
| 52 | + self.output_dir = Path(output_dir) |
| 53 | + self.verbose = verbose |
| 54 | + self.use_cache = use_cache |
| 55 | + self.version = version |
| 56 | + |
| 57 | + if self.verbose: |
| 58 | + logger.setLevel(logging.DEBUG) |
| 59 | + |
| 60 | + # Configure retry strategy with httpx |
| 61 | + transport = httpx.HTTPTransport( |
| 62 | + retries=3, |
| 63 | + ) |
| 64 | + |
| 65 | + self.client = httpx.Client( |
| 66 | + timeout=30.0, |
| 67 | + transport=transport, |
| 68 | + headers=self.headers, |
| 69 | + follow_redirects=True, |
| 70 | + ) |
| 71 | + |
| 72 | + def get_contents(self, path=""): |
| 73 | + url = f"{self.base_url}/contents/{path}" |
| 74 | + logger.info(f"Fetching contents from: {url}") |
| 75 | + |
| 76 | + try: |
| 77 | + response = self.client.get(url) |
| 78 | + response.raise_for_status() |
| 79 | + |
| 80 | + # Rate limiting info |
| 81 | + remaining = response.headers.get("X-RateLimit-Remaining") |
| 82 | + if remaining: |
| 83 | + logger.info(f"API calls remaining: {remaining}") |
| 84 | + |
| 85 | + content = response.json() |
| 86 | + if isinstance(content, str): |
| 87 | + raise ValueError(f"Unexpected response format: {content}") |
| 88 | + return content |
| 89 | + |
| 90 | + except HTTPStatusError as e: |
| 91 | + logger.error(f"HTTP error occurred: {e}") |
| 92 | + raise |
| 93 | + except RequestError as e: |
| 94 | + logger.error(f"Request error occurred: {e}") |
| 95 | + raise |
| 96 | + |
| 97 | + def load_cached_specs(self) -> dict[str, list[dict[str, str]]]: |
| 98 | + try: |
| 99 | + if self.cache_file.exists() and self.cache_file.stat().st_size > 0: |
| 100 | + with open(self.cache_file) as f: |
| 101 | + cached = json.load(f) |
| 102 | + return cached |
| 103 | + except (json.JSONDecodeError, OSError) as e: |
| 104 | + logger.warning(f"Cache read error: {e}") |
| 105 | + return {"": []} |
| 106 | + |
| 107 | + def save_specs_cache(self, specs): |
| 108 | + try: |
| 109 | + with open(self.cache_file, "w") as f: |
| 110 | + json.dump(specs, f, indent=2) |
| 111 | + except OSError as e: |
| 112 | + logger.error(f"Failed to save cache: {e}") |
| 113 | + |
| 114 | + def discover_specs(self, use_cache=True) -> dict[str, list[dict[str, str]]]: |
| 115 | + if use_cache: |
| 116 | + cached = self.load_cached_specs() |
| 117 | + if cached: |
| 118 | + logger.info("Using cached specs") |
| 119 | + return cached |
| 120 | + |
| 121 | + specs = {self.version: []} |
| 122 | + specs_for_version = specs[self.version] |
| 123 | + |
| 124 | + # Check core path |
| 125 | + core_specs = self.get_contents("core") |
| 126 | + for spec in core_specs: |
| 127 | + if spec.get("type") == "file" and spec.get("name", "").endswith(".json"): |
| 128 | + specs_for_version.append( |
| 129 | + { |
| 130 | + "name": Path(spec["name"]).stem, |
| 131 | + "url": f"{self.raw_base}/core/{spec['name']}", |
| 132 | + } |
| 133 | + ) |
| 134 | + |
| 135 | + # Check apps path |
| 136 | + apps = self.get_contents("apps") |
| 137 | + for app in apps: |
| 138 | + if app.get("type") == "dir": |
| 139 | + versions = self.get_contents(app["path"]) |
| 140 | + for version in versions: |
| 141 | + if version.get("type") == "dir": |
| 142 | + files = self.get_contents(version["path"]) |
| 143 | + for file in files: |
| 144 | + if file.get("name", "").endswith(".json"): |
| 145 | + specs_for_version.append( |
| 146 | + { |
| 147 | + "name": app["name"].split(".")[0], |
| 148 | + "url": f"{self.raw_base}/{file['path']}", |
| 149 | + } |
| 150 | + ) |
| 151 | + |
| 152 | + if specs: |
| 153 | + self.save_specs_cache(specs) |
| 154 | + |
| 155 | + return specs |
| 156 | + |
| 157 | + def generate_models(self): |
| 158 | + output_dir = self.output_dir |
| 159 | + output_dir.mkdir(exist_ok=True) |
| 160 | + specs: dict[str, list[dict[str, str]]] = self.discover_specs( |
| 161 | + use_cache=self.use_cache |
| 162 | + ) |
| 163 | + |
| 164 | + if not specs: |
| 165 | + logger.warning("No specs found!") |
| 166 | + return |
| 167 | + |
| 168 | + if self.version not in specs: |
| 169 | + logger.warning(f"No specs found for version {self.version}") |
| 170 | + return |
| 171 | + |
| 172 | + for spec in specs[self.version]: |
| 173 | + url_parts = spec["url"].split("/") |
| 174 | + module_name = url_parts[-1].replace(".json", "") |
| 175 | + |
| 176 | + cmd = [ |
| 177 | + "datamodel-codegen", |
| 178 | + "--url", |
| 179 | + spec["url"], |
| 180 | + "--output-model-type", |
| 181 | + "pydantic_v2.BaseModel", |
| 182 | + "--use-annotated", |
| 183 | + "--enum-field-as-literal", |
| 184 | + "all", |
| 185 | + "--output", |
| 186 | + str(output_dir), |
| 187 | + ] |
| 188 | + |
| 189 | + # the core API should use the file name as the output of the DMCG command |
| 190 | + # the core app does not have apps dir in its URL |
| 191 | + if "apps" not in url_parts and module_name == "core": |
| 192 | + cmd[-1] = str(output_dir) + "/core.py" |
| 193 | + |
| 194 | + try: |
| 195 | + logger.info(f"Generating models for {module_name}...") |
| 196 | + subprocess.run(cmd, check=True) |
| 197 | + except subprocess.CalledProcessError as e: |
| 198 | + logger.error(f"Error generating models for {module_name}: {e}") |
| 199 | + |
| 200 | + def __del__(self): |
| 201 | + self.client.close() |
| 202 | + |
| 203 | + |
| 204 | +if __name__ == "__main__": |
| 205 | + parser = argparse.ArgumentParser( |
| 206 | + description="Discover OpenAPI specifications and generate Pydantic models." |
| 207 | + ) |
| 208 | + parser.add_argument( |
| 209 | + "--input", |
| 210 | + type=str, |
| 211 | + default="./openapi_specs", |
| 212 | + help="Path to the OpenAPI specifications directory. Default: ./openapi_specs", |
| 213 | + ) |
| 214 | + parser.add_argument( |
| 215 | + "--output", |
| 216 | + type=str, |
| 217 | + default="./pydantic_eda", |
| 218 | + help="Path to the output directory.", |
| 219 | + ) |
| 220 | + parser.add_argument( |
| 221 | + "--version", |
| 222 | + type=str, |
| 223 | + default="main", |
| 224 | + help="openapi repo version (tag) to get the models from. Default: main", |
| 225 | + ) |
| 226 | + parser.add_argument( |
| 227 | + "--verbose", |
| 228 | + action="store_true", |
| 229 | + help="Enable verbose logging. Default: False", |
| 230 | + ) |
| 231 | + parser.add_argument( |
| 232 | + "--no-cache", |
| 233 | + action="store_true", |
| 234 | + help="Force fresh discovery by ignoring the cache. Default: False", |
| 235 | + ) |
| 236 | + |
| 237 | + args = parser.parse_args() |
| 238 | + |
| 239 | + discovery = OpenAPIDiscovery( |
| 240 | + input_dir=args.input, |
| 241 | + output_dir=args.output, |
| 242 | + version=args.version, |
| 243 | + verbose=args.verbose, |
| 244 | + use_cache=not args.no_cache, |
| 245 | + ) |
| 246 | + discovery.generate_models() |
0 commit comments