Merge pull request #91 from ponytailer/swagger-cli

Add swagger cli

Author: ponytailer

.github/workflows/ci.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
       fail-fast: true
 
     steps:

poetry.lock

@@ -0,0 +1,190 @@
+import logging
+from collections import defaultdict
+from typing import Optional
+
+from openapi_spec_validator import validate
+from openapi_spec_validator.readers import read_from_filename
+from pydantic import BaseModel
+from typer import Typer
+
+all_models = {}
+
+logger = logging.getLogger(__name__)
+
+cli_app = Typer()
+
+
+class RequestEntity(BaseModel):
+    method: str
+    function_name: str
+    path: str
+    response: str
+    arg_dict: dict
+    args: Optional[str] = ""
+
+
+def generate_pydantic_model(schema, model_name):
+    properties = schema.get('properties', {})
+    required_fields = schema.get('required', [])
+
+    simple_fields = {}
+    for field_name, field_schema in properties.items():
+        field_type = field_schema.get('type')
+        if field_type == 'object':
+            ref_object = field_schema["$ref"].rsplit("/", 1)[-1]
+            simple_fields[field_name] = ref_object
+        elif field_type == 'array':
+            items_schema = field_schema.get('items', {})
+            items_type = items_schema.get('type')
+            if items_type == 'object':
+                ref_object = field_schema["$ref"].rsplit("/", 1)[-1]
+                simple_fields[field_name] = list[ref_object]
+            else:
+                simple_fields[field_name] = list[
+                    map_type_to_string(items_type)]
+        else:
+            simple_fields[field_name] = map_type_to_string(
+                field_type)
+
+        if field_name not in required_fields:
+            simple_fields[field_name] = f"Optional[{simple_fields[field_name]}]"
+
+    if simple_fields:
+        all_models[model_name] = simple_fields
+
+
+def map_type_to_string(openapi_type):
+    type_mapping = {
+        'string': "str",
+        'integer': "int",
+        'number': "float",
+        'boolean': "bool",
+        'array': list,
+        'object': dict
+    }
+    return type_mapping.get(openapi_type, str)  # 默认返回 str
+
+
+def success_return(status, response):
+    return status == "200" or status == "201" and "content" in response
+
+
+def convert_request_name(name: str):
+    s = []
+    for k in name:
+        if k.isupper():
+            s.append("_")
+            s.append(k.lower())
+        else:
+            s.append(k)
+    return "".join(s[1:]) if s[0] == "_" else "".join(s)
+
+
+def parse_swagger_and_generate_models(swagger_path):
+    swagger_dict, _ = read_from_filename(swagger_path)
+    validate(swagger_dict)
+
+    definitions = swagger_dict.get('components', {}).get('schemas', {})
+    paths = swagger_dict.get('paths', {})
+
+    for model_name, model_schema in definitions.items():
+        generate_pydantic_model(model_schema, model_name)
+
+    path_pairs = defaultdict(list)
+
+    for path, path_item in paths.items():
+        function_name_suffix = "_".join(
+            p for p in path.split("/") if "{" not in p)
+
+        for method, operation in path_item.items():
+            request_entity_body = {}
+            arg_dict = {}
+
+            # 处理路径参数和查询参数
+            if "parameters" in operation:
+                for parameter in operation["parameters"]:
+                    if parameter["in"] == "path":
+                        for _, value_ in parameter["schema"].items():
+                            value_string = map_type_to_string(value_)
+                            arg_dict[parameter["name"]] = value_string
+                            break
+                    elif parameter["in"] == "query":
+                        for _, value_ in parameter["schema"].items():
+                            value_string = map_type_to_string(value_)
+                            arg_dict[
+                                parameter["name"]] = f"Optional[{value_string}]"
+                            break
+                    else:
+                        logger.warning(
+                            f"unknown parameter in {path}: {parameter}")
+
+            # 处理请求体
+            if "requestBody" in operation:
+                request_body = operation['requestBody']
+                for _, media_type in request_body.get("content", {}).items():
+                    schema = media_type.get("schema", {})
+                    if not schema:
+                        continue
+                    request_model_name = schema.get("$ref", "").split('/')[-1]
+                    arg_dict[convert_request_name(
+                        request_model_name)] = request_model_name
+                    break
+
+            # 处理响应体
+            if 'responses' in operation:
+                for status_code, response in operation['responses'].items():
+                    if not success_return(status_code, response):
+                        continue
+                    for _, media_type in response["content"].items():
+                        if "schema" not in media_type:
+                            continue
+                        schema = media_type['schema']
+                        # 提取响应模型名称
+                        response_model_name = \
+                            schema.get('$ref', '').split('/')[-1]
+                        request_entity_body["response"] = response_model_name
+
+            request_entity_body["method"] = method
+            request_entity_body[
+                "function_name"] = f"{method}{function_name_suffix}"
+            request_entity_body["arg_dict"] = arg_dict
+            path_pairs[path].append(request_entity_body)
+
+    ret = []
+    for path, values in path_pairs.items():
+        for value in values:
+            value["path"] = path
+            value["function_name"] = value["function_name"].replace("-", "_")
+            ret.append(RequestEntity(**value))
+    return ret
+
+
+@cli_app.command()
+def parse(path: str, model_file_name: str = ""):
+    import jinja2
+    results = parse_swagger_and_generate_models(path)
+    tmpl = jinja2.Template(open("pydantic_client/models.template").read())
+    render_fields = {
+        key: [
+            f"{field}: {ftype}"
+            for field, ftype in value.items()
+        ] for key, value in all_models.items()
+    }
+    for entity in results:
+        args = ", ".join([
+            f"{name}: {type_}"
+            for name, type_ in entity.arg_dict.items()
+        ])
+        entity.args = args
+
+    if render_fields:
+        model_string = tmpl.render(models=render_fields, info=results)
+        if model_file_name:
+            with open(model_file_name, "w") as f:
+                f.write(model_string)
+        else:
+            print(model_string)
+
+
+if __name__ == "__main__":
+    cli_app()

pydantic_client/cli.py

@@ -0,0 +1,16 @@
+from typing import Optional
+from pydantic import BaseModel
+{% for model_name, fields in models.items() %}
+
+class {{model_name}}(BaseModel):
+    {% for field in fields -%}
+    {{ field }}
+    {% endfor %}
+{% endfor %}
+
+class WebClient:
+    {% for entity in info %}
+    @{{entity.method}}("{{entity.path}}")
+    def {{entity.function_name}}(self{%if entity.args%}, {%endif%}{{entity.args}}):
+        ...
+    {% endfor %}

pydantic_client/models.template

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pydantic-client"
-version = "1.0.6"
+version = "1.0.7"
 description = "Http client base pydantic, with requests or aiohttp"
 authors = ["ponytailer <huangxiaohen2738@gmail.com>"]
 readme = "README.md"
@@ -13,19 +13,25 @@ url = "https://pypi.tuna.tsinghua.edu.cn/simple"
 
 
 [tool.poetry.dependencies]
-python = "^3.8"
+python = "^3.9"
 
 pydantic = ">=2.1"
 requests = "*"
 
 aiohttp = { version = "*", optional = true }
 httpx = { version = "*", extras = ["http2"], optional = true }
+openapi-spec-validator = "^0.7.1"
+jinja2 = "^3.1.5"
+typer = "^0.15.1"
 
 [tool.poetry.extras]
 httpx = ["httpx"]
 aiohttp = ["aiohttp"]
 all = ["httpx", "aiohttp"]
 
+[tool.poetry.scripts]
+pydantic-client = "pydantic_client.cli:cli_app"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

pyproject.toml

@@ -17,4 +17,7 @@ pydantic==2.5.2
 python-multipart==0.0.9
 requests==2.31.0
 uvicorn==0.29.0
-toml
+toml
+typer
+jinja2
+openapi-spec-validator