Add type annotations to models (#1200)

* Add type annotations to models

This adds explicit parameters for optional paramaters and annotations
for all of them, making sure oneOf are included.
This also fixes oneOf handling in API body parameters.

* Remove positional arguments

* Remove the args

* Don't use oneOf object

* Push nullable everywhere

* pre-commit fixes

* Fix oneOf parameter import

* Black

Co-authored-by: ci.datadog-api-spec <[email protected]>

Author: therve

.generator/src/generator/cli.py

@@ -42,9 +42,11 @@ def cli(specs, output):
     env.globals["package"] = PACKAGE_NAME
     env.globals["get_name"] = formatter.get_name
     env.globals["get_type_for_attribute"] = openapi.get_type_for_attribute
+    env.globals["get_typing_for_attribute"] = openapi.get_typing_for_attribute
     env.globals["get_types_for_attribute"] = openapi.get_types_for_attribute
     env.globals["get_type_for_parameter"] = openapi.get_type_for_parameter
     env.globals["get_references_for_model"] = openapi.get_references_for_model
+    env.globals["get_oneof_references_for_model"] = openapi.get_oneof_references_for_model
     env.globals["get_oneof_parameters"] = openapi.get_oneof_parameters
     env.globals["get_type_for_items"] = openapi.get_type_for_items
     env.globals["get_api_models"] = openapi.get_api_models

.generator/src/generator/formatter.py

@@ -15,10 +15,7 @@
 PRIMITIVE_TYPES = ["string", "number", "boolean", "integer"]
 
 EDGE_CASES = {}
-replacement_file = (
-    pathlib.Path(__file__).parent
-    / "replacement.json"
-)
+replacement_file = pathlib.Path(__file__).parent / "replacement.json"
 if replacement_file.exists():
     with replacement_file.open() as f:
         EDGE_CASES.update(json.load(f))
@@ -96,7 +93,12 @@ def header(self, text, level, raw=None):
 
 
 def docstring(text):
-    return m2r2.convert(text.replace("\\n", "\\\\n"), renderer=CustomRenderer())[1:-1].replace("\\ ", " ").replace("\\`", "\\\\`").replace("\n\n\n", "\n\n")
+    return (
+        m2r2.convert(text.replace("\\n", "\\\\n"), renderer=CustomRenderer())[1:-1]
+        .replace("\\ ", " ")
+        .replace("\\`", "\\\\`")
+        .replace("\n\n\n", "\n\n")
+    )
 
 
 def _merge_imports(a, b):
@@ -111,10 +113,7 @@ def format_parameters(kwargs, spec, version, replace_values=None):
     imports = defaultdict(set)
 
     parameters_spec = {p["name"]: p for p in spec.get("parameters", [])}
-    if (
-        "requestBody" in spec
-        and "multipart/form-data" in spec["requestBody"]["content"]
-    ):
+    if "requestBody" in spec and "multipart/form-data" in spec["requestBody"]["content"]:
         parent = spec["requestBody"]["content"]["multipart/form-data"]["schema"]
         for name, schema in parent["properties"].items():
             parameters_spec[name] = {
@@ -153,9 +152,7 @@ def get_name_and_imports(schema, version=None, imports=None):
         name = schema.__reference__["$ref"].split("/")[-1]
         if "oneOf" not in schema:
             # do not include parent of oneOf schema
-            imports[
-                MODEL_IMPORT_TPL.format(version=version, name=safe_snake_case(name))
-            ].add(name)
+            imports[MODEL_IMPORT_TPL.format(version=version, name=safe_snake_case(name))].add(name)
 
     return name, imports
 
@@ -176,10 +173,10 @@ def format_data_with_schema(
     imports = imports or defaultdict(set)
     if schema.get("type") not in {"string", "integer", "boolean"} or schema.get("enum"):
         name, imports = get_name_and_imports(schema, version, imports)
+    if schema.get("oneOf"):
+        name = None
     if name:
-        imports[
-            MODEL_IMPORT_TPL.format(version=version, name=safe_snake_case(name))
-        ].add(name)
+        imports[MODEL_IMPORT_TPL.format(version=version, name=safe_snake_case(name))].add(name)
 
     if "enum" in schema and data not in schema["enum"]:
         raise ValueError(f"{data} is not valid enum value {schema['enum']}")
@@ -321,15 +318,9 @@ def format_data_with_schema_dict(
             imports = _merge_imports(imports, extra_imports)
 
     if not name and "oneOf" not in schema:
-        if (
-            default_name
-            and not schema.get("additionalProperties")
-            and schema.get("properties")
-        ):
+        if default_name and not schema.get("additionalProperties") and schema.get("properties"):
             name = default_name
-            imports[
-                MODEL_IMPORT_TPL.format(version=version, name=safe_snake_case(name))
-            ].add(name)
+            imports[MODEL_IMPORT_TPL.format(version=version, name=safe_snake_case(name))].add(name)
         else:
             name = "dict"
             warnings.warn(f"Unnamed schema {schema} for {data}")
@@ -360,7 +351,7 @@ def format_data_with_schema_dict(
             warnings.warn(f"[{matched}] {data} is not valid for schema {name}")
 
     if not parameters and data:
-        parameters = ", ".join(f"{k}=\"{v}\"" for k, v in data.items())
+        parameters = ", ".join(f'{k}="{v}"' for k, v in data.items())
 
     if name:
         return f"{name}({parameters})", imports