🐛 Fix: webhook parsing when having nested union types (#57)

Co-authored-by: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>

codegen/__init__.py

@@ -9,7 +9,7 @@
 from .config import Config
 from .source import get_source
 from .log import logger as logger
-from .parser.schemas import ModelSchema
+from .parser.schemas import ModelSchema, UnionSchema
 from .parser import (
     WebhookData,
     EndpointData,
@@ -26,24 +26,17 @@
     lstrip_blocks=True,
     extensions=["jinja2.ext.loopcontrols"],
 )
-env.globals.update(
-    {
-        "repr": repr,
-        "sanitize": sanitize,
-        "snake_case": snake_case,
-        "pascal_case": pascal_case,
-        "kebab_case": kebab_case,
-    }
-)
-env.filters.update(
-    {
-        "repr": repr,
-        "sanitize": sanitize,
-        "snake_case": snake_case,
-        "pascal_case": pascal_case,
-        "kebab_case": kebab_case,
-    }
-)
+
+_funcs = {
+    "repr": repr,
+    "sanitize": sanitize,
+    "snake_case": snake_case,
+    "pascal_case": pascal_case,
+    "kebab_case": kebab_case,
+    "is_union_schema": lambda x: isinstance(x, UnionSchema),
+}
+env.globals.update(_funcs)
+env.filters.update(_funcs)
 
 
 def load_config() -> Config:
@@ -198,9 +191,7 @@ def build_legacy_rest_models(
     logger.info("Successfully generated legacy rest models!")
 
 
-def build_versions(
-    dir: Path, output_module: str, versions: dict[str, str], latest_version: str
-):
+def build_versions(dir: Path, versions: dict[str, str], latest_version: str):
     logger.info("Start generating versions...")
 
     # build __init__.py
@@ -216,23 +207,15 @@ def build_versions(
     rest_template = env.get_template("versions/rest.py.jinja")
     rest_path = dir / "rest.py"
     rest_path.write_text(
-        rest_template.render(
-            output_module=output_module,
-            versions=versions,
-            latest_version=latest_version,
-        )
+        rest_template.render(versions=versions, latest_version=latest_version)
     )
 
     # build webhooks.py
     logger.info("Building versions webhooks.py...")
     webhooks_template = env.get_template("versions/webhooks.py.jinja")
     webhooks_path = dir / "webhooks.py"
     webhooks_path.write_text(
-        webhooks_template.render(
-            output_module=output_module,
-            versions=versions,
-            latest_version=latest_version,
-        )
+        webhooks_template.render(versions=versions, latest_version=latest_version)
     )
 
     logger.info("Successfully generated versions!")
@@ -332,7 +315,7 @@ def build():
         latest_model_names,
         latest_event_names,
     )
-    build_versions(config.output_dir, output_module, versions, latest_version)
+    build_versions(config.output_dir, versions, latest_version)
     build_legacy_rest_models(
         config.legacy_rest_models,
         output_module,

codegen/templates/versions/rest.py.jinja

@@ -45,7 +45,7 @@ class RestVersionSwitcher(_VersionProxy):
     def __call__(self, version: VERSION_TYPE = LATEST_VERSION) -> Any:
         if version in self._cached_namespaces:
             return self._cached_namespaces[version]
-        module = importlib.import_module(f"{{ output_module }}.{VERSIONS[version]}.rest", __name__)
+        module = importlib.import_module(f".{VERSIONS[version]}.rest", __package__)
         namespace = module.RestNamespace(self._github)
         self._cached_namespaces[version] = namespace
         return namespace

codegen/templates/versions/webhooks.py.jinja

@@ -41,7 +41,7 @@ class WebhooksVersionSwitcher(_VersionProxy):
     def __call__(self, version: VERSION_TYPE = LATEST_VERSION) -> Any:
         if version in self._cached_namespaces:
             return self._cached_namespaces[version]
-        module = importlib.import_module(f"{{ output_module }}.{VERSIONS[version]}.webhooks", __name__)
+        module = importlib.import_module(f".{VERSIONS[version]}.webhooks", __package__)
         namespace = module.WebhookNamespace()
         self._cached_namespaces[version] = namespace
         return namespace

codegen/templates/webhooks/_namespace.py.jinja

@@ -66,7 +66,7 @@ class WebhookNamespace:
         """
         if name not in VALID_EVENT_NAMES:
             raise WebhookTypeNotFound(name)
-        module = importlib.import_module(f".{name}", __name__)
+        module = importlib.import_module(f".{name}", __package__)
         Event = getattr(module, "Event")
         return type_validate_json(Event, payload)
 
@@ -106,7 +106,7 @@ class WebhookNamespace:
 
         if name not in VALID_EVENT_NAMES:
             raise WebhookTypeNotFound(name)
-        module = importlib.import_module(f".{name}", __name__)
+        module = importlib.import_module(f".{name}", __package__)
         Event = getattr(module, "Event")
         return type_validate_python(Event, payload)
 

codegen/templates/webhooks/event.py.jinja

@@ -7,6 +7,7 @@ from typing_extensions import TypeAlias, Annotated
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 {% for webhook in webhooks %}
@@ -18,7 +19,18 @@ from githubkit.compat import GitHubModel
 {% if webhooks | length > 1 %}
 Event: TypeAlias = Annotated[Union[
     {% for webhook in webhooks %}
+    {% if is_union_schema(webhook.event_schema) %}
+    Annotated[
+        {{ webhook.event_schema.get_type_string() }},
+        TaggedUnion(
+            {{ webhook.event_schema.get_type_string() }},
+            "action",
+            "{{ webhook.action }}"
+        )
+    ],
+    {% else %}
     {{ webhook.event_schema.get_type_string() }},
+    {% endif %}
     {% endfor %}
 ], Field(discriminator="action")]
 {% else %}

githubkit/lazy_module.py

@@ -9,11 +9,11 @@
 from importlib.machinery import ModuleSpec, PathFinder, SourceFileLoader
 
 LAZY_MODULES = (
-    r"githubkit\.rest",
-    r"githubkit\.versions\.v[^.]+\.webhooks",
-    r"githubkit\.versions\.latest\.models",
-    r"githubkit\.versions\.latest\.types",
-    r"githubkit\.versions\.latest\.webhooks",
+    r"^githubkit\.rest$",
+    r"^githubkit\.versions\.v[^.]+\.webhooks$",
+    r"^githubkit\.versions\.latest\.models$",
+    r"^githubkit\.versions\.latest\.types$",
+    r"^githubkit\.versions\.latest\.webhooks$",
 )
 
 

githubkit/utils.py

@@ -1,10 +1,15 @@
 import inspect
 from enum import Enum
-from typing import Any, Dict, Literal, TypeVar, final
+from typing import Any, Dict, Type, Generic, Literal, TypeVar, final
 
+from pydantic import BaseModel
 from pydantic_core import to_jsonable_python
 
-from .compat import custom_validation
+from .compat import PYDANTIC_V2, custom_validation, type_validate_python
+
+if PYDANTIC_V2:
+    from pydantic_core import core_schema
+    from pydantic import GetCoreSchemaHandler
 
 T = TypeVar("T")
 
@@ -75,3 +80,34 @@ def obj_to_jsonable(obj: Any) -> Any:
         return obj
 
     return to_jsonable_python(obj)
+
+
+class TaggedUnion(Generic[T]):
+    __slots__ = ("type_", "discriminator", "tag")
+
+    def __init__(self, type_: Type[T], discriminator: str, tag: str) -> None:
+        self.type_ = type_
+        self.discriminator = discriminator
+        self.tag = tag
+
+    def _validate(self, value: Any) -> T:
+        return type_validate_python(self.type_, value)
+
+    if PYDANTIC_V2:
+
+        def __get_pydantic_core_schema__(
+            self, _source_type: Any, _handler: "GetCoreSchemaHandler"
+        ) -> "core_schema.CoreSchema":
+            return core_schema.no_info_before_validator_function(
+                self._validate,
+                core_schema.model_schema(
+                    BaseModel,
+                    schema=core_schema.model_fields_schema(
+                        {
+                            self.discriminator: core_schema.model_field(
+                                core_schema.literal_schema([self.tag])
+                            )
+                        }
+                    ),
+                ),
+            )

githubkit/versions/rest.py

@@ -49,9 +49,7 @@ def __call__(self) -> "V20221128RestNamespace":
     def __call__(self, version: VERSION_TYPE = LATEST_VERSION) -> Any:
         if version in self._cached_namespaces:
             return self._cached_namespaces[version]
-        module = importlib.import_module(
-            f"githubkit.versions.{VERSIONS[version]}.rest", __name__
-        )
+        module = importlib.import_module(f".{VERSIONS[version]}.rest", __package__)
         namespace = module.RestNamespace(self._github)
         self._cached_namespaces[version] = namespace
         return namespace

githubkit/versions/v2022_11_28/webhooks/_namespace.py

@@ -731,7 +731,7 @@ def parse(name: EventNameType, payload: Union[str, bytes]) -> "WebhookEvent":
         """
         if name not in VALID_EVENT_NAMES:
             raise WebhookTypeNotFound(name)
-        module = importlib.import_module(f".{name}", __name__)
+        module = importlib.import_module(f".{name}", __package__)
         Event = getattr(module, "Event")
         return type_validate_json(Event, payload)
 
@@ -1223,7 +1223,7 @@ def parse_obj(name: EventNameType, payload: Dict[str, Any]) -> "WebhookEvent":
 
         if name not in VALID_EVENT_NAMES:
             raise WebhookTypeNotFound(name)
-        module = importlib.import_module(f".{name}", __name__)
+        module = importlib.import_module(f".{name}", __package__)
         Event = getattr(module, "Event")
         return type_validate_python(Event, payload)
 

githubkit/versions/v2022_11_28/webhooks/branch_protection_configuration.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import (

githubkit/versions/v2022_11_28/webhooks/branch_protection_rule.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import (

githubkit/versions/v2022_11_28/webhooks/check_run.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import (

githubkit/versions/v2022_11_28/webhooks/check_suite.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import (

githubkit/versions/v2022_11_28/webhooks/code_scanning_alert.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import (

githubkit/versions/v2022_11_28/webhooks/commit_comment.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import WebhookCommitCommentCreated

githubkit/versions/v2022_11_28/webhooks/create.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import WebhookCreate

githubkit/versions/v2022_11_28/webhooks/custom_property.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import (

githubkit/versions/v2022_11_28/webhooks/custom_property_values.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import WebhookCustomPropertyValuesUpdated

githubkit/versions/v2022_11_28/webhooks/delete.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import WebhookDelete

githubkit/versions/v2022_11_28/webhooks/dependabot_alert.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import (

githubkit/versions/v2022_11_28/webhooks/deploy_key.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import WebhookDeployKeyCreated, WebhookDeployKeyDeleted

githubkit/versions/v2022_11_28/webhooks/deployment.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import WebhookDeploymentCreated

githubkit/versions/v2022_11_28/webhooks/deployment_protection_rule.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import WebhookDeploymentProtectionRuleRequested

githubkit/versions/v2022_11_28/webhooks/deployment_review.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import (

githubkit/versions/v2022_11_28/webhooks/deployment_status.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import WebhookDeploymentStatusCreated

githubkit/versions/v2022_11_28/webhooks/discussion.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import (

githubkit/versions/v2022_11_28/webhooks/discussion_comment.py

@@ -12,6 +12,7 @@
 
 from pydantic import Field
 
+from githubkit.utils import TaggedUnion
 from githubkit.compat import GitHubModel
 
 from ..models import (