🚸 improve parser and code generator

Author: yanyongyu

codegen/__init__.py

@@ -40,15 +40,15 @@ def build_templates(data: GeneratorData, config: Config):
     models_template = env.get_template("models/models.py.jinja")
     models_path = Path(config.models_output)
     models_path.parent.mkdir(parents=True, exist_ok=True)
-    models_path.write_text(models_template.render(data=data))
+    models_path.write_text(models_template.render(models=data.models))
     logger.info("Successfully built models!")
 
     # build types
     logger.info("Building types...")
-    types_template = env.get_template("types/types.py.jinja")
+    types_template = env.get_template("models/types.py.jinja")
     types_path = Path(config.types_output)
     types_path.parent.mkdir(parents=True, exist_ok=True)
-    types_path.write_text(types_template.render(data=data))
+    types_path.write_text(types_template.render(models=data.models))
     logger.info("Successfully built types!")
 
     # build endpoints

codegen/parser/data.py

@@ -1,12 +1,10 @@
-from itertools import chain
 from collections import defaultdict
-from typing import Set, Dict, List, Optional
+from typing import Dict, List, Optional
 
 from pydantic import BaseModel
 
 from .endpoints import EndpointData
 from .schemas import SchemaData, ModelSchema
-from .utils import snake_case, fix_reserved_words
 
 
 class GeneratorData(BaseModel):
@@ -28,6 +26,3 @@ def endpoints_by_tag(self) -> Dict[str, List[EndpointData]]:
     @property
     def models(self) -> List[ModelSchema]:
         return [schema for schema in self.schemas if isinstance(schema, ModelSchema)]
-
-    def get_imports(self) -> Set[str]:
-        return set(chain.from_iterable(schema.get_imports() for schema in self.schemas))

codegen/parser/endpoints/request_body.py

@@ -31,12 +31,10 @@ def get_params_defination(self) -> List[str]:
 
     def get_param_imports(self) -> Set[str]:
         imports = set()
+        imports.update(self.body_schema.get_param_imports())
         if isinstance(self.body_schema, ModelSchema):
-            imports.update(self.body_schema.get_param_imports())
             for prop in self.body_schema.properties:
                 imports.update(prop.get_param_imports())
-        else:
-            imports.update(self.body_schema.get_param_imports())
         return imports
 
     def get_using_imports(self) -> Set[str]:

codegen/parser/schemas/schema.py

@@ -20,19 +20,23 @@ def get_param_type_string(self) -> str:
         """Get type string used by client codegen"""
         return self.get_type_string()
 
-    def get_imports(self) -> Set[str]:
-        """Get schema needed imports for creating the schema"""
+    def get_model_imports(self) -> Set[str]:
+        """Get schema needed imports for model codegen"""
         return set()
 
+    def get_type_imports(self) -> Set[str]:
+        """Get schema needed imports for types codegen"""
+        return self.get_model_imports()
+
     def get_param_imports(self) -> Set[str]:
-        """Get schema needed imports for typing params"""
-        return self.get_imports()
+        """Get schema needed imports for client param codegen"""
+        return self.get_model_imports()
 
     def get_using_imports(self) -> Set[str]:
-        """Get schema needed imports for using"""
-        return self.get_imports()
+        """Get schema needed imports for client request codegen"""
+        return self.get_model_imports()
 
-    def get_default_args(self) -> Dict[str, str]:
+    def _get_default_args(self) -> Dict[str, str]:
         """Get pydantic field info args"""
         default = self.default
         args = {}
@@ -59,18 +63,26 @@ def get_type_string(self) -> str:
         return type_string if self.required else f"Union[Unset, {type_string}]"
 
     def get_param_type_string(self) -> str:
+        """Get type string used by client codegen"""
         type_string = self.schema_data.get_param_type_string()
         return type_string if self.required else f"Union[Unset, {type_string}]"
 
-    def get_imports(self) -> Set[str]:
-        """Get schema needed imports for creating the schema"""
-        imports = self.schema_data.get_imports()
+    def get_model_imports(self) -> Set[str]:
+        """Get schema needed imports for model codegen"""
+        imports = self.schema_data.get_model_imports()
         imports.add("from pydantic import Field")
         if not self.required:
             imports.add("from typing import Union")
             imports.add("from githubkit.utils import UNSET, Unset")
         return imports
 
+    def get_type_imports(self) -> Set[str]:
+        """Get schema needed imports for type codegen"""
+        imports = self.schema_data.get_type_imports()
+        if not self.required:
+            imports.add("from typing_extensions import NotRequired")
+        return imports
+
     def get_param_imports(self) -> Set[str]:
         """Get schema needed imports for typing params"""
         imports = self.schema_data.get_param_imports()
@@ -79,8 +91,8 @@ def get_param_imports(self) -> Set[str]:
             imports.add("from githubkit.utils import UNSET, Unset")
         return imports
 
-    def get_default_string(self) -> str:
-        args = self.schema_data.get_default_args()
+    def _get_default_string(self) -> str:
+        args = self.schema_data._get_default_args()
         if "default" not in args and "default_factory" not in args:
             args["default"] = "..." if self.required else "UNSET"
         if self.prop_name != self.name:
@@ -103,12 +115,12 @@ def get_param_defination(self) -> str:
     def get_model_defination(self) -> str:
         """Get defination used by model codegen"""
         type_ = self.get_type_string()
-        default = self.get_default_string()
+        default = self._get_default_string()
         return f"{self.prop_name}: {type_} = {default}"
 
     def get_type_defination(self) -> str:
         """Get defination used by types codegen"""
-        type_ = self.get_param_type_string()
+        type_ = self.schema_data.get_param_type_string()
         return (
             f"{self.prop_name}: {type_ if self.required else f'NotRequired[{type_}]'}"
         )
@@ -117,8 +129,8 @@ def get_type_defination(self) -> str:
 class AnySchema(SchemaData):
     _type_string: ClassVar[str] = "Any"
 
-    def get_imports(self) -> Set[str]:
-        imports = super().get_imports()
+    def get_model_imports(self) -> Set[str]:
+        imports = super().get_model_imports()
         imports.add("from typing import Any")
         return imports
 
@@ -140,8 +152,8 @@ class IntSchema(SchemaData):
 
     _type_string: ClassVar[str] = "int"
 
-    def get_default_args(self) -> Dict[str, str]:
-        args = super().get_default_args()
+    def _get_default_args(self) -> Dict[str, str]:
+        args = super()._get_default_args()
         if self.multiple_of is not None:
             args["multiple_of"] = repr(self.multiple_of)
         if self.maximum is not None:
@@ -164,8 +176,8 @@ class FloatSchema(SchemaData):
 
     _type_string: ClassVar[str] = "float"
 
-    def get_default_args(self) -> Dict[str, str]:
-        args = super().get_default_args()
+    def _get_default_args(self) -> Dict[str, str]:
+        args = super()._get_default_args()
         if self.multiple_of is not None:
             args["multiple_of"] = str(self.multiple_of)
         if self.maximum is not None:
@@ -186,8 +198,8 @@ class StringSchema(SchemaData):
 
     _type_string: ClassVar[str] = "str"
 
-    def get_default_args(self) -> Dict[str, str]:
-        args = super().get_default_args()
+    def _get_default_args(self) -> Dict[str, str]:
+        args = super()._get_default_args()
         if self.min_length is not None:
             args["min_length"] = str(self.min_length)
         if self.max_length is not None:
@@ -200,26 +212,26 @@ def get_default_args(self) -> Dict[str, str]:
 class DateTimeSchema(SchemaData):
     _type_string: ClassVar[str] = "datetime"
 
-    def get_imports(self) -> Set[str]:
-        imports = super().get_imports()
+    def get_model_imports(self) -> Set[str]:
+        imports = super().get_model_imports()
         imports.add("from datetime import datetime")
         return imports
 
 
 class DateSchema(SchemaData):
     _type_string: ClassVar[str] = "date"
 
-    def get_imports(self) -> Set[str]:
-        imports = super().get_imports()
+    def get_model_imports(self) -> Set[str]:
+        imports = super().get_model_imports()
         imports.add("from datetime import date")
         return imports
 
 
 class FileSchema(SchemaData):
     _type_string: ClassVar[str] = "FileTypes"
 
-    def get_imports(self) -> Set[str]:
-        imports = super().get_imports()
+    def get_model_imports(self) -> Set[str]:
+        imports = super().get_model_imports()
         imports.add("from githubkit.typing import FileTypes")
         return imports
 
@@ -236,10 +248,15 @@ def get_type_string(self) -> str:
     def get_param_type_string(self) -> str:
         return f"List[{self.item_schema.get_param_type_string()}]"
 
-    def get_imports(self) -> Set[str]:
-        imports = super().get_imports()
+    def get_model_imports(self) -> Set[str]:
+        imports = super().get_model_imports()
         imports.add("from typing import List")
-        imports.update(self.item_schema.get_imports())
+        imports.update(self.item_schema.get_model_imports())
+        return imports
+
+    def get_type_imports(self) -> Set[str]:
+        imports = {"from typing import List"}
+        imports.update(self.item_schema.get_type_imports())
         return imports
 
     def get_param_imports(self) -> Set[str]:
@@ -252,8 +269,8 @@ def get_using_imports(self) -> Set[str]:
         imports.update(self.item_schema.get_using_imports())
         return imports
 
-    def get_default_args(self) -> Dict[str, str]:
-        args = super().get_default_args()
+    def _get_default_args(self) -> Dict[str, str]:
+        args = super()._get_default_args()
         # FIXME: remove list constraints due to forwardref not supported
         # See https://github.com/samuelcolvin/pydantic/issues/3745
         if isinstance(self.item_schema, (ModelSchema, UnionSchema)):
@@ -282,8 +299,8 @@ def is_float_enum(self) -> bool:
     def get_type_string(self) -> str:
         return f"Literal[{', '.join(repr(value) for value in self.values)}]"
 
-    def get_imports(self) -> Set[str]:
-        imports = super().get_imports()
+    def get_model_imports(self) -> Set[str]:
+        imports = super().get_model_imports()
         imports.add("from typing import Literal")
         return imports
 
@@ -299,13 +316,19 @@ def get_type_string(self) -> str:
     def get_param_type_string(self) -> str:
         return f"{self.class_name}Type"
 
-    def get_imports(self) -> Set[str]:
-        imports = super().get_imports()
+    def get_model_imports(self) -> Set[str]:
+        imports = super().get_model_imports()
         imports.add("from pydantic import BaseModel")
         if self.allow_extra:
             imports.add("from pydantic import Extra")
         for prop in self.properties:
-            imports.update(prop.get_imports())
+            imports.update(prop.get_model_imports())
+        return imports
+
+    def get_type_imports(self) -> Set[str]:
+        imports = {"from typing_extensions import TypedDict"}
+        for prop in self.properties:
+            imports.update(prop.get_type_imports())
         return imports
 
     def get_param_imports(self) -> Set[str]:
@@ -331,11 +354,17 @@ def get_param_type_string(self) -> str:
             return self.schemas[0].get_param_type_string()
         return f"Union[{', '.join(schema.get_param_type_string() for schema in self.schemas)}]"
 
-    def get_imports(self) -> Set[str]:
-        imports = super().get_imports()
+    def get_model_imports(self) -> Set[str]:
+        imports = super().get_model_imports()
         imports.add("from typing import Union")
         for schema in self.schemas:
-            imports.update(schema.get_imports())
+            imports.update(schema.get_model_imports())
+        return imports
+
+    def get_type_imports(self) -> Set[str]:
+        imports = {"from typing import Union"}
+        for schema in self.schemas:
+            imports.update(schema.get_type_imports())
         return imports
 
     def get_param_imports(self) -> Set[str]:
@@ -350,11 +379,11 @@ def get_using_imports(self) -> Set[str]:
             imports.update(schema.get_using_imports())
         return imports
 
-    def get_default_args(self) -> Dict[str, str]:
+    def _get_default_args(self) -> Dict[str, str]:
         args = {}
         for schema in self.schemas:
-            args.update(schema.get_default_args())
-        args.update(super().get_default_args())
+            args.update(schema._get_default_args())
+        args.update(super()._get_default_args())
         if self.discriminator:
             args["discriminator"] = self.discriminator
         return args

codegen/templates/client/param.py.jinja renamed to codegen/templates/client/_param.py.jinja

@@ -1,4 +1,3 @@
-
 {% macro path_params(endpoint) %}
 {% for path_param in endpoint.path_params %}
 {{ path_param.get_param_defination() }},

codegen/templates/client/request.py.jinja renamed to codegen/templates/client/_request.py.jinja

@@ -1,4 +1,4 @@
-{% from "client/response.py.jinja" import build_response_model, build_error_models %}
+{% from "client/_response.py.jinja" import build_response_model, build_error_models %}
 
 {% macro build_path(endpoint) %}
 {% if endpoint.path_params %}

codegen/templates/client/response.py.jinja renamed to codegen/templates/client/_response.py.jinja

@@ -4,9 +4,9 @@ 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/response.py.jinja" import build_response_type %}
-{% from "client/request.py.jinja" import build_request, build_request_params %}
+{% from "client/_param.py.jinja" import endpoint_params %}
+{% from "client/_response.py.jinja" import build_response_type %}
+{% from "client/_request.py.jinja" import build_request, build_request_params %}
 
 {% for endpoint in endpoints %}
 {% for import_ in endpoint.get_imports() %}

codegen/templates/client/client.py.jinja

@@ -0,0 +1,15 @@
+{% macro build_model_docstring(model) %}
+"""{{ model.title or model.class_name | wordwrap(80) }}
+{% if model.description %}
+
+{{ model.description | wordwrap(80) }}
+{% endif %}
+{% if model.examples %}
+
+Examples:
+    {% for example in model.examples %}
+    {{ example | string | wordwrap(82) }}
+    {% endfor %}
+{% endif %}
+"""
+{% endmacro %}