✨ add overloads for rest api

Author: yanyongyu

codegen/parser/endpoints/endpoint.py

@@ -59,6 +59,9 @@ def get_imports(self) -> Set[str]:
         if self.request_body:
             imports.update(self.request_body.get_param_imports())
             imports.update(self.request_body.get_using_imports())
+            if self.request_body.allowed_models:
+                imports.add("from typing import Union")
+                imports.add("from githubkit.utils import UNSET, Unset")
         if self.success_response:
             imports.update(self.success_response.get_using_imports())
         for resp in self.error_responses.values():

codegen/parser/endpoints/request_body.py

@@ -5,7 +5,7 @@
 
 from ...source import Source
 from ..utils import concat_snake_name
-from ..schemas import Property, SchemaData, ModelSchema, parse_schema
+from ..schemas import Property, SchemaData, ModelSchema, UnionSchema, parse_schema
 
 
 class RequestBodyData(BaseModel):
@@ -14,26 +14,40 @@ class RequestBodyData(BaseModel):
     required: bool = False
 
     @property
-    def is_model(self) -> bool:
-        return isinstance(self.body_schema, ModelSchema)
+    def allowed_models(self) -> List[ModelSchema]:
+        if isinstance(self.body_schema, ModelSchema):
+            return [self.body_schema]
+        elif isinstance(self.body_schema, UnionSchema):
+            return [
+                schema
+                for schema in self.body_schema.schemas
+                if isinstance(schema, ModelSchema)
+            ]
+        return []
 
-    def get_params_defination(self) -> List[str]:
-        if not isinstance(self.body_schema, ModelSchema):
-            prop = Property(
-                name="body",
-                prop_name="body",
-                required=self.required,
-                schema_data=self.body_schema,
-            )
-            return [prop.get_param_defination()]
+    def get_raw_definition(self) -> str:
+        prop = Property(
+            name="data",
+            prop_name="data",
+            required=self.required,
+            schema_data=self.body_schema,
+        )
+        return prop.get_param_defination()
 
-        return [prop.get_param_defination() for prop in self.body_schema.properties]
+    def get_endpoint_definition(self) -> str:
+        prop = Property(
+            name="data",
+            prop_name="data",
+            required=not bool(self.allowed_models),
+            schema_data=self.body_schema,
+        )
+        return prop.get_param_defination()
 
     def get_param_imports(self) -> Set[str]:
         imports = set()
         imports.update(self.body_schema.get_param_imports())
-        if isinstance(self.body_schema, ModelSchema):
-            for prop in self.body_schema.properties:
+        for model in self.allowed_models:
+            for prop in model.properties:
                 imports.update(prop.get_param_imports())
         return imports
 

codegen/templates/client/_param.py.jinja

@@ -22,19 +22,39 @@
 {% endfor %}
 {% endmacro %}
 
-{% macro body_params(endpoint) %}
-{% for prop in endpoint.request_body.get_params_defination() %}
-{{ prop }},
+{% macro body_params(model) %}
+{% for prop in model.properties %}
+{{ prop.get_param_defination() }},
 {% endfor %}
 {% endmacro %}
 
-{% macro endpoint_params(endpoint) %}
+{% macro endpoint_raw_params(endpoint) %}
+{{ path_params(endpoint) }}
+{{ query_params(endpoint) }}
+{{ header_params(endpoint) }}
+{{ cookie_params(endpoint) }}
+*,
+{{ endpoint.request_body.get_raw_definition() }}
+{% endmacro %}
+
+{% macro endpoint_model_params(endpoint, model) %}
+{{ path_params(endpoint) }}
+{{ query_params(endpoint) }}
+{{ header_params(endpoint) }}
+{{ cookie_params(endpoint) }}
+*,
+data: Unset = UNSET,
+{{ body_params(model) }}
+{% endmacro %}
+
+{% macro endpoint_params(endpoint, model) %}
 {{ path_params(endpoint) }}
 {{ query_params(endpoint) }}
 {{ header_params(endpoint) }}
 {{ cookie_params(endpoint) }}
 {%- if endpoint.request_body %}
 *,
-{{ body_params(endpoint) }}
+{{ endpoint.request_body.get_endpoint_definition() }},
+**kwargs
 {%- endif %}
 {% endmacro %}

codegen/templates/client/_request.py.jinja

@@ -1,5 +1,7 @@
 {% from "client/_response.py.jinja" import build_response_model, build_error_models %}
 
+{% set TYPE_MAPPING = {"json": "json", "form": "data", "file": "files", "raw": "content"} %}
+
 {% macro build_path(endpoint) %}
 {% if endpoint.path_params %}
 url = f"{{ endpoint.path }}"
@@ -32,28 +34,17 @@ cookies = {
 }
 {% endmacro %}
 
-{% macro _get_body_name(request_body) %}
-{% if request_body.is_model %}
-{{ request_body.body_schema.class_name }}(**{
-    {% for prop in request_body.body_schema.properties %}
-    "{{ prop.name }}": {{ prop.prop_name }},
-    {% endfor %}
-}).dict(by_alias=True)
-{% else %}
-body
-{% endif %}
-{% endmacro %}
-
 {% macro build_body(request_body) %}
-{% if request_body.type == "json" %}
-json = {{ _get_body_name(request_body) }}
-{% elif request_body.type == "form" %}
-data = {{ _get_body_name(request_body) }}
-{% elif request_body.type == "file" %}
-files = {{ _get_body_name(request_body) }}
-{% elif request_body.type == "raw" %}
-content = {{ _get_body_name(request_body) }}
-{% endif %}
+{% set name = TYPE_MAPPING[request_body.type] %}
+if not kwargs:
+    kwargs = UNSET
+
+{{ name }} = kwargs if data is UNSET else data
+{{ name }} = parse_obj_as(
+    {{ request_body.body_schema.get_type_string() }},
+    {{ name }}
+)
+{{ name }} = {{ name }}.dict(by_alias=True) if isinstance({{ name }}, BaseModel) else {{ name }}
 {% endmacro %}
 
 {% macro build_request(endpoint) %}
@@ -73,15 +64,8 @@ url,
 params=exclude_unset(params),
 {% endif %}
 {% if endpoint.request_body %}
-{% if endpoint.request_body.type == "raw" %}
-content=exclude_unset(content),
-{% elif endpoint.request_body.type == "form" %}
-data=exclude_unset(data),
-{% elif endpoint.request_body.type == "file" %}
-files=exclude_unset(files),
-{% elif endpoint.request_body.type == "json" %}
-json=exclude_unset(json),
-{% endif %}
+{% set name = TYPE_MAPPING[endpoint.request_body.type] %}
+{{ name }}=exclude_unset({{ name }}),
 {% endif %}
 {% if endpoint.header_params %}
 headers=exclude_unset(headers),

codegen/templates/client/client.py.jinja

@@ -4,7 +4,7 @@ This file is auto generated by github rest api discription.
 See https://github.com/github/rest-api-description for more information.
 """
 
-{% from "client/_param.py.jinja" import endpoint_params %}
+{% from "client/_param.py.jinja" import endpoint_params, endpoint_raw_params, endpoint_model_params %}
 {% from "client/_response.py.jinja" import build_response_type %}
 {% from "client/_request.py.jinja" import build_request, build_request_params %}
 
@@ -14,7 +14,9 @@ See https://github.com/github/rest-api-description for more information.
 {% endfor %}
 {% endfor %}
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
+
+from pydantic import BaseModel, parse_obj_as
 
 from githubkit.utils import exclude_unset
 
@@ -27,6 +29,26 @@ class {{ pascal_case(tag) }}Client:
         self._github = github
 
     {% for endpoint in endpoints %}
+    {% if endpoint.request_body and endpoint.request_body.allowed_models %}
+    {# generate raw data overload #}
+    @overload
+    def {{ endpoint.name }}(
+        self,
+        {{ endpoint_raw_params(endpoint) | indent(8) }}
+    ) -> "{{ build_response_type(endpoint.success_response) }}":
+        ...
+
+    {# generate model data overload #}
+    {% for model in endpoint.request_body.allowed_models %}
+    @overload
+    def {{ endpoint.name }}(
+        self,
+        {{ endpoint_model_params(endpoint, model) | indent(8) }}
+    ) -> "{{ build_response_type(endpoint.success_response) }}":
+        ...
+
+    {% endfor %}
+    {% endif %}
     def {{ endpoint.name }}(
         self,
         {{ endpoint_params(endpoint) | indent(8) }}
@@ -36,6 +58,26 @@ class {{ pascal_case(tag) }}Client:
             {{ build_request_params(endpoint) | indent(12) }}
         )
 
+    {% if endpoint.request_body and endpoint.request_body.allowed_models %}
+    {# generate raw data overload #}
+    @overload
+    async def async_{{ endpoint.name }}(
+        self,
+        {{ endpoint_raw_params(endpoint) | indent(8) }}
+    ) -> "{{ build_response_type(endpoint.success_response) }}":
+        ...
+
+    {# generate model data overload #}
+    {% for model in endpoint.request_body.allowed_models %}
+    @overload
+    async def async_{{ endpoint.name }}(
+        self,
+        {{ endpoint_model_params(endpoint, model) | indent(8) }}
+    ) -> "{{ build_response_type(endpoint.success_response) }}":
+        ...
+
+    {% endfor %}
+    {% endif %}
     async def async_{{ endpoint.name }}(
         self,
         {{ endpoint_params(endpoint) | indent(8) }}

codegen/templates/models/models.py.jinja

@@ -15,7 +15,7 @@ from __future__ import annotations
 {% endfor %}
 
 
-class GitHubModel(BaseModel, arbitrary_types_allowed=True):
+class GitHubModel(BaseModel, allow_population_by_field_name=True):
     ...
 
 {# model #}

codegen/templates/namespace/namespace.py.jinja

@@ -17,8 +17,6 @@ if TYPE_CHECKING:
 
 
 class RestNamespace:
-    __slots__ = ("_github",)
-
     def __init__(self, github: "GitHubCore"):
         self._github = github
 

githubkit/rest/__init__.py

@@ -50,8 +50,6 @@
 
 
 class RestNamespace:
-    __slots__ = ("_github",)
-
     def __init__(self, github: "GitHubCore"):
         self._github = github