client code generation

Author: FaydSpeare

README.md

@@ -1,7 +1,7 @@
 # EnsembleData Python API
 
 [![pypi](https://img.shields.io/pypi/v/ensembledata?color=%2334D058&label=pypi%20package)](https://pypi.org/project/ensembledata/)
-![](https://img.shields.io/pypi/pyversions/ensembledata.svg?color=%2334D058)
+[![pypi](https://img.shields.io/pypi/pyversions/ensembledata.svg)](https://pypi.org/project/ensembledata/)
 
 ## Documentation
 
@@ -28,7 +28,7 @@ from ensembledata.api import EDClient
 
 
 client = EDClient("API-TOKEN")
-result = client.tiktok.user_info_from_username("daviddobrik")
+result = client.tiktok.user_info_from_username(username="daviddobrik")
 
 print("Data: ", result.data)
 print("Units charged:", result.units_charged)
@@ -79,7 +79,7 @@ from ensembledata.api import EDClient, EDError, errors
 
 client = EDClient("API-TOKEN")
 try:
-    result = client.tiktok.user_info_from_username("daviddobrik")
+    result = client.tiktok.user_info_from_username(username="daviddobrik")
 except EDError as e:
 
     # Rate limit exceeded...
@@ -113,7 +113,7 @@ from ensembledata.api import EDAsyncClient
 
 async def main():
     client = EDAsyncClient("API-TOKEN")
-    result = await client.tiktok.user_info_from_username("daviddobrik")
+    result = await client.tiktok.user_info_from_username(username="daviddobrik")
 
 if __name__ == "__main__":
     asyncio.run(main())

codegen/__init__.py

@@ -0,0 +1,101 @@
+from typing import Any
+
+import requests
+from jinja2 import Environment, PackageLoader, select_autoescape
+
+
+def openapi_type_to_python_type(openapi_type: str) -> str:
+    if openapi_type == "integer":
+        return "int"
+    if openapi_type == "string":
+        return "str"
+    if openapi_type == "boolean":
+        return "bool"
+    raise Exception("Unknown type")
+
+
+def parse_endpoints_by_tag(openapi: Any):
+    tags = dict()
+    for path, endpoints in openapi["paths"].items():
+        operation_id = endpoints["get"]["operationId"]
+        function_name = "_".join(operation_id.split("_")[1:])
+        params = []
+        tag = endpoints["get"]["tags"][0]
+
+        for param in endpoints["get"]["parameters"]:
+            name = param["name"]
+            if name == "token":
+                continue
+
+            type = None
+            transform = None
+
+            if "type" in param["schema"]:
+                type = openapi_type_to_python_type(param["schema"]["type"])
+
+                if "retype" in param["schema"]:  # noqa: SIM102
+                    if param["schema"]["retype"] == "semicolon-separated-string-to-list":
+                        type = "Sequence[str]"
+                        transform = "join_with_semicolon"
+
+            # Handle enums
+            elif "allOf" in param["schema"]:
+                assert len(param["schema"]["allOf"]) == 1
+                assert "$ref" in param["schema"]["allOf"][0]
+                ref = param["schema"]["allOf"][0]["$ref"]
+                ref_name = ref.split("/")[-1]
+                type = f"Literal{openapi['components']['schemas'][ref_name]['enum']}"
+            else:
+                raise Exception("Unknown param type")
+
+            params.append(
+                {
+                    "required": param["required"],
+                    "param": name,
+                    "name": param["schema"].get("rename") or name,
+                    "type": type,
+                    "transform": transform,
+                }
+            )
+
+        if tag not in tags:
+            tags[tag] = []
+
+        tags[tag].append(
+            {
+                "function_name": function_name,
+                "path": path,
+                "params": params,
+                "return_top_level_data": operation_id
+                in ["tiktok_user_posts_from_username", "tiktok_user_posts_from_secuid"],
+            }
+        )
+
+    return tags
+
+
+def main():
+    openapi = requests.get("https://ensembledata.com/apis/openapi.json").json()
+
+    env = Environment(
+        loader=PackageLoader("codegen"),
+        autoescape=select_autoescape(),
+        lstrip_blocks=True,
+        trim_blocks=True,
+    )
+
+    tags = parse_endpoints_by_tag(openapi)
+    client_template = env.get_template("client.jinja")
+
+    client_content = client_template.render(tags=tags, async_methods=False)
+    async_client_content = client_template.render(tags=tags, async_methods=True)
+
+    with open("ensembledata/api/_client.py", "w") as file:
+        file.write(client_content)
+
+    with open("ensembledata/api/_async_client.py", "w") as file:
+        file.write(async_client_content)
+
+
+if __name__ == "__main__":
+    main()

codegen/__main__.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING, Any, Mapping, Sequence
+
+if sys.version_info < (3, 8):
+    from typing_extensions import Literal
+else:
+    from typing import Literal
+
+from ._requester import {{ "Async" if async_methods else "" }}Requester
+
+if TYPE_CHECKING:
+    from ._response import EDResponse
+
+class UseDefault:
+    pass
+
+USE_DEFAULT = UseDefault()
+
+{% for key, endpoints in tags.items() %}
+class {{ key }}Endpoints:
+
+    def __init__(self, requester: {{ "Async" if async_methods else "" }}Requester):
+        self.requester = requester
+
+    {% for x in endpoints %}
+    {{ "async " if async_methods else "" }}def {{ x.function_name }}(
+        self,
+        *,
+        {% for item in x.params %}
+        {% if item.required %}
+        {{ item.name }}: {{ item.type }},
+        {% else %}
+        {{ item.name }}: {{ item.type }} | UseDefault = USE_DEFAULT,
+        {% endif %}
+        {% endfor %}        
+        extra_params: Mapping[str, Any] | None = None,
+    ) -> EDResponse:
+        params: dict[str, Any] = {
+        {% for item in x.params %}
+        {% if item.transform == "join_with_semicolon" %}
+            "{{ item.param }}": ";".join({{ item.name }}),
+        {% else %}
+            "{{ item.param }}": {{ item.name }},
+        {% endif %}
+        {% endfor %}
+        }
+        if extra_params is not None:
+            params = {**extra_params, **params}
+        params = {k: v for k, v in params.items() if not (v is None or v is USE_DEFAULT)}
+        return {{ "await " if async_methods else "" }}self.requester.get("{{ x.path }}", params=params{{ ", return_top_level_data=True" if x.return_top_level_data else "" }}) 
+    
+    {% endfor %}
+
+{% endfor %}
+
+class ED{{ "Async" if async_methods else "" }}Client:
+    def __init__(self, token: str, *, timeout: int = 600, max_network_retries: int = 3):
+        self.requester = {{ "Async" if async_methods else "" }}Requester(
+            token, timeout=timeout, max_network_retries=max_network_retries
+        )
+        {% for tag in tags %}
+        self.{{ tag|lower }} = {{ tag }}Endpoints(self.requester)
+        {% endfor %}
+
+    async def request(self, uri: str, params: Mapping[str, Any] | None = None) -> EDResponse:
+        return {{ "await " if async_methods else "" }}self.requester.get(uri, params=params or {})

codegen/templates/client.jinja

@@ -1,8 +1,8 @@
 from ._async_client import EDAsyncClient
 from ._client import EDClient
 from ._response import EDResponse
-from .errors import EDError
 from ._version import version
+from .errors import EDError
 
 __version__ = version
 __all__ = ["EDClient", "EDAsyncClient", "EDResponse", "EDError"]