♻️ refactor the codegen for future

Author: yanyongyu

codegen/__init__.py

@@ -7,8 +7,8 @@
 from jinja2 import Environment, PackageLoader
 
 from .log import logger
-from .config import Config
 from .source import get_source
+from .config import Config, RestConfig, WebhookConfig
 from .parser import (
     OpenAPIData,
     WebhookData,
@@ -39,16 +39,14 @@
 def load_config() -> Config:
     pyproject = tomli.loads(Path("./pyproject.toml").read_text())
     config_dict: Dict[str, Any] = pyproject.get("tool", {}).get("codegen", {})
-    config_dict = {
-        k.replace("--", "").replace("-", "_"): v for k, v in config_dict.items()
-    }
+
     return Config.parse_obj(config_dict)
 
 
-def build_rest_api(data: OpenAPIData, config: Config):
+def build_rest_api(data: OpenAPIData, rest: RestConfig):
     logger.info("Start generating rest api codes...")
 
-    client_path = Path(config.client_output)
+    client_path = Path(rest.output_dir)
     shutil.rmtree(client_path)
     client_path.mkdir(parents=True, exist_ok=True)
 
@@ -74,7 +72,7 @@ def build_rest_api(data: OpenAPIData, config: Config):
         tag_path = client_path / f"{tag}.py"
         tag_path.write_text(
             client_template.render(
-                tag=tag, endpoints=endpoints, rest_api_version=config.rest_api_version
+                tag=tag, endpoints=endpoints, rest_api_version=rest.version
             )
         )
         logger.info(f"Successfully built endpoints for tag {tag}!")
@@ -92,21 +90,21 @@ def build_rest_api(data: OpenAPIData, config: Config):
     logger.info("Successfully generated rest api codes!")
 
 
-def build_webhook(data: WebhookData, config: Config):
+def build_webhook(data: WebhookData, webhook: WebhookConfig):
     logger.info("Start generating webhook codes...")
 
     # build models
     logger.info("Building webhook models...")
     models_template = env.get_template("models/webhooks.py.jinja")
-    models_path = Path(config.webhooks_output)
+    models_path = Path(webhook.output)
     models_path.parent.mkdir(parents=True, exist_ok=True)
     models_path.write_text(models_template.render(models=data.models))
     logger.info("Successfully built webhook models!")
 
     # build types
     logger.info("Building webhook types...")
     types_template = env.get_template("models/webhook_types.py.jinja")
-    types_path = Path(config.webhook_types_output)
+    types_path = Path(webhook.types_output)
     types_path.parent.mkdir(parents=True, exist_ok=True)
     types_path.write_text(
         types_template.render(
@@ -133,30 +131,31 @@ def build():
     config = load_config()
     logger.info(f"Loaded config: {config!r}")
 
-    logger.info("Start getting OpenAPI source...")
-    source = get_source(httpx.URL(config.rest_description_source))
-    logger.info(f"Getting schema from {source.uri} succeeded!")
-
-    logger.info("Start parsing OpenAPI spec...")
-    _patch_openapi_spec(source.root)
-    parsed_data = parse_openapi_spec(source, config)
-    logger.info(
-        "Successfully parsed OpenAPI spec: "
-        f"{len(parsed_data.schemas)} schemas, {len(parsed_data.endpoints)} endpoints"
-    )
+    for versioned_rest in config.rest:
+        logger.info(f"Start getting OpenAPI source for {versioned_rest.version}...")
+        source = get_source(httpx.URL(versioned_rest.description_source))
+        logger.info(f"Getting schema from {source.uri} succeeded!")
+
+        logger.info(f"Start parsing OpenAPI spec for {versioned_rest.version}...")
+        _patch_openapi_spec(source.root)
+        parsed_data = parse_openapi_spec(source, versioned_rest, config)
+        logger.info(
+            f"Successfully parsed OpenAPI spec {versioned_rest.version}: "
+            f"{len(parsed_data.schemas)} schemas, {len(parsed_data.endpoints)} endpoints"
+        )
 
-    build_rest_api(parsed_data, config)
+        build_rest_api(parsed_data, versioned_rest)
 
-    del parsed_data
+        del parsed_data
 
     logger.info("Start getting Webhook source...")
-    source = get_source(httpx.URL(config.webhook_schema_source))
+    source = get_source(httpx.URL(config.webhook.schema_source))
     logger.info(f"Getting schema from {source.uri} succeeded!")
 
     logger.info("Start parsing Webhook spec...")
-    parsed_data = parse_webhook_schema(source, config)
+    parsed_data = parse_webhook_schema(source, config.webhook, config)
     logger.info(
         f"Successfully parsed Webhook spec: {len(parsed_data.definitions)} schemas"
     )
 
-    build_webhook(parsed_data, config)
+    build_webhook(parsed_data, config.webhook)

codegen/config.py

@@ -1,17 +1,27 @@
-from typing import Any, Dict
+from typing import Any, Dict, List
 
-from pydantic import BaseModel
 import openapi_schema_pydantic as oas
+from pydantic import Field, BaseModel
 
 
-class Config(BaseModel):
-    rest_description_source: str
-    rest_api_version: str
-    webhook_schema_source: str
-    class_overrides: Dict[str, str] = {}
-    field_overrides: Dict[str, str] = {}
-    schema_overrides: Dict[str, Dict[str, Any]] = {}
+class Overridable(BaseModel):
+    class_overrides: Dict[str, str] = Field(default_factory=dict)
+    field_overrides: Dict[str, str] = Field(default_factory=dict)
+    schema_overrides: Dict[str, Dict[str, Any]] = Field(default_factory=dict)
 
-    client_output: str
-    webhooks_output: str
-    webhook_types_output: str
+
+class RestConfig(Overridable):
+    version: str
+    description_source: str
+    output_dir: str
+
+
+class WebhookConfig(Overridable):
+    schema_source: str
+    output: str
+    types_output: str
+
+
+class Config(Overridable):
+    rest: List[RestConfig]
+    webhook: WebhookConfig

codegen/parser/__init__.py

@@ -1,16 +1,16 @@
 from contextvars import ContextVar
-from typing import Dict, List, Union, Optional
+from typing import Dict, List, Tuple, Union, Optional
 
 import httpx
 from openapi_schema_pydantic import OpenAPI
 
 # parser context
+_override_config: ContextVar[Tuple["Overridable", ...]] = ContextVar("override_config")
 _schemas: ContextVar[Dict[httpx.URL, "SchemaData"]] = ContextVar("schemas")
-_config: ContextVar["Config"] = ContextVar("config")
 
 
-def get_config() -> "Config":
-    return _config.get()
+def get_override_config() -> Tuple["Overridable", ...]:
+    return _override_config.get()
 
 
 def get_schemas() -> Dict[httpx.URL, "SchemaData"]:
@@ -27,7 +27,6 @@ def add_schema(ref: httpx.URL, schema: "SchemaData"):
     _schemas.get()[ref] = schema
 
 
-from ..config import Config
 from ..source import Source
 from .endpoints import parse_endpoint
 from .utils import sanitize as sanitize
@@ -39,13 +38,14 @@ def add_schema(ref: httpx.URL, schema: "SchemaData"):
 from .endpoints import EndpointData as EndpointData
 from .schemas import SchemaData, UnionSchema, parse_schema
 from .utils import fix_reserved_words as fix_reserved_words
+from ..config import Config, RestConfig, Overridable, WebhookConfig
 
 
-def parse_openapi_spec(source: Source, config: Config) -> OpenAPIData:
+def parse_openapi_spec(source: Source, rest: RestConfig, config: Config) -> OpenAPIData:
     source = source.get_root()
 
+    _ot = _override_config.set((rest, config))
     _st = _schemas.set({})
-    _ct = _config.set(config)
 
     try:
         openapi = OpenAPI.parse_obj(source.root)
@@ -70,15 +70,17 @@ def parse_openapi_spec(source: Source, config: Config) -> OpenAPIData:
             schemas=list(get_schemas().values()),
         )
     finally:
+        _override_config.reset(_ot)
         _schemas.reset(_st)
-        _config.reset(_ct)
 
 
-def parse_webhook_schema(source: Source, config: Config) -> WebhookData:
+def parse_webhook_schema(
+    source: Source, webhook: WebhookConfig, config: Config
+) -> WebhookData:
     source = source.get_root()
 
+    _ot = _override_config.set((webhook, config))
     _st = _schemas.set({})
-    _ct = _config.set(config)
 
     try:
         root_schema = parse_schema(source, "webhook_schema")
@@ -106,5 +108,5 @@ def parse_webhook_schema(source: Source, config: Config) -> WebhookData:
             definitions=definitions,
         )
     finally:
+        _override_config.reset(_ot)
         _schemas.reset(_st)
-        _config.reset(_ct)

codegen/parser/utils.py

@@ -7,8 +7,8 @@
 from pydantic import parse_obj_as
 import openapi_schema_pydantic as oas
 
-from . import get_config
 from ..source import Source
+from . import get_override_config
 
 DELIMITERS = r"\. _-"
 
@@ -80,20 +80,29 @@ def build_boolean(value: Union[bool, str]) -> bool:
 
 
 def build_class_name(name: str) -> str:
-    config = get_config()
+    sources = get_override_config()
     class_name = fix_reserved_words(pascal_case(name))
-    return config.class_overrides.get(class_name, class_name)
+    for override_source in sources:
+        if override := override_source.class_overrides.get(class_name):
+            return override
+    return class_name
 
 
 def build_prop_name(name: str) -> str:
-    config = get_config()
-    name = config.field_overrides.get(name, name)
+    sources = get_override_config()
+    for override_source in sources:
+        if override := override_source.field_overrides.get(name):
+            name = override
+            break
     return fix_reserved_words(snake_case(name))
 
 
 def get_schema_override(source: Source) -> Optional[Dict[str, Any]]:
-    config = get_config()
-    return config.schema_overrides.get(source.pointer.path, None)
+    sources = get_override_config()
+    for override_source in sources:
+        if schema := override_source.schema_overrides.get(source.pointer.path):
+            return schema
+    return None
 
 
 def merge_dict(old: dict, new: dict):

codegen/source.py

@@ -38,7 +38,7 @@ def __truediv__(self, other: Union[str, int]) -> "Source":
 
 
 @cache
-def get_content(source: Union[httpx.URL, Path]) -> OpenAPI:
+def get_content(source: Union[httpx.URL, Path]) -> dict:
     return (
         json.loads(source.read_text())
         if isinstance(source, Path)

codegen/templates/client/_request.py.jinja

@@ -11,27 +11,32 @@ url = "{{ endpoint.path }}"
 {% endmacro %}
 
 {% macro build_query(params) %}
+{% if params %}
 params = {
     {% for param in params %}
     "{{ param.name }}": {{ param.prop_name }},
     {% endfor %}
 }
+{% endif %}
 {% endmacro %}
 
 {% macro build_header(params) %}
 headers = {
     {% for param in params %}
     "{{ param.name }}": {{ param.prop_name }},
     {% endfor %}
+    "X-GitHub-Api-Version": self._REST_API_VERSION,
 }
 {% endmacro %}
 
 {% macro build_cookie(params) %}
+{% if params %}
 cookies = {
     {% for param in params %}
     "{{ param.name }}": {{ param.prop_name }},
     {% endfor %}
 }
+{% endif %}
 {% endmacro %}
 
 {% macro build_body(request_body) %}
@@ -49,9 +54,9 @@ if not kwargs:
 
 {% macro build_request(endpoint) %}
 {{ build_path(endpoint) }}
-{% if endpoint.query_params %}
 {{ build_query(endpoint.query_params) }}
-{% endif %}
+{{ build_header(endpoint.header_params) }}
+{{ build_cookie(endpoint.cookie_params) }}
 {% if endpoint.request_body %}
 {{ build_body(endpoint.request_body) }}
 {% endif %}
@@ -67,11 +72,7 @@ params=exclude_unset(params),
 {% set name = TYPE_MAPPING[endpoint.request_body.type] %}
 {{ name }}=exclude_unset({{ name }}),
 {% endif %}
-{% if endpoint.header_params %}
-headers=exclude_unset({**headers, "X-GitHub-Api-Version": self._REST_API_VERSION}),
-{% else %}
-headers={"X-GitHub-Api-Version": self._REST_API_VERSION},
-{% endif %}
+headers=exclude_unset(headers),
 {% if endpoint.cookie_params %}
 cookies=exclude_unset(cookies),
 {% endif %}