Cleanup formatter module (#1026)

And fixes a bug with keywords for examples.

Author: therve

.generator/conftest.py

@@ -15,13 +15,8 @@
 from pytest_bdd import given, parsers, then, when
 
 from generator import openapi
-from generator.utils import camel_case, given_variables, snake_case, untitle_case
 
-from generator.python_formatter import (
-    format_parameters as format_parameters_python,
-    format_data_with_schema as format_data_with_schema_python,
-    safe_snake_case,
-)
+from generator.formatter import format_parameters, format_data_with_schema, safe_snake_case, snake_case
 
 
 MODIFIED_FEATURES = {
@@ -84,28 +79,18 @@ def encode(self, obj):
         return result
 
 
-def json_dumps(*args, **kwargs):
-    if CLIENT_REPO_NAME == "terraform-config":
-        kwargs.setdefault("cls", FloatEncoder)
-    return json.dumps(*args, **kwargs)
-
-
 JINJA_ENV = Environment(loader=FileSystemLoader(pathlib.Path(__file__).parent / "src" / "generator" / "templates"))
 JINJA_ENV.filters["tojson"] = json.dumps
 JINJA_ENV.filters["snake_case"] = snake_case
 JINJA_ENV.filters["safe_snake_case"] = safe_snake_case
-JINJA_ENV.filters["camel_case"] = camel_case
-JINJA_ENV.filters["untitle_case"] = untitle_case
-JINJA_ENV.globals["format_data_with_schema_python"] = format_data_with_schema_python
-JINJA_ENV.globals["format_parameters_python"] = format_parameters_python
-JINJA_ENV.globals["given_variables"] = given_variables
+JINJA_ENV.globals["format_data_with_schema"] = format_data_with_schema
+JINJA_ENV.globals["format_parameters"] = format_parameters
 
 PYTHON_EXAMPLE_J2 = JINJA_ENV.get_template("example.j2")
 
 
 def pytest_bdd_after_scenario(request, feature, scenario):
     try:
-        specs = request.getfixturevalue("specs")
         operation_specs = request.getfixturevalue("operation_specs")
         version = request.getfixturevalue("api_version")
         context = request.getfixturevalue("context")
@@ -547,7 +532,7 @@ def the_status_is(context, status, description):
 @then(parsers.parse('the response "{response_path}" is equal to {value}'))
 def expect_equal(context, response_path, value):
     """Compare a response attribute to a value."""
-    pass
+
 
 @then(
     parsers.parse(
@@ -556,16 +541,13 @@ def expect_equal(context, response_path, value):
 )
 def expect_equal_value(context, response_path, fixture_path):
     """Compare a response attribute to another attribute."""
-    pass
 
 
 @then(parsers.parse('the response "{response_path}" has length {fixture_length:d}'))
 def expect_equal_length(context, response_path, fixture_length):
     """Check the length of a response attribute."""
-    pass
 
 
 @then(parsers.parse('the response "{response_path}" is false'))
 def expect_false(context, response_path):
     """Check that a response attribute is false."""
-    pass

.generator/src/generator/formatter.py

@@ -1,9 +1,27 @@
 """Data formatter."""
+from collections import defaultdict
+import json
+from functools import singledispatch
+import pathlib
 import keyword
+import warnings
 import re
 
+import dateutil.parser
 import m2r2
 
+
+MODEL_IMPORT_TPL = "datadog_api_client.{version}.model.{name}"
+
+EDGE_CASES = {}
+replacement_file = (
+    pathlib.Path(__file__).parent
+    / "replacement.json"
+)
+if replacement_file.exists():
+    with replacement_file.open() as f:
+        EDGE_CASES.update(json.load(f))
+
 KEYWORDS = set(keyword.kwlist)
 KEYWORDS.add("property")
 
@@ -21,6 +39,12 @@ def snake_case(value):
     return PATTERN_DOUBLE_UNDERSCORE.sub("_", s1)
 
 
+def safe_snake_case(value):
+    for token, replacement in EDGE_CASES.items():
+        value = value.replace(token, replacement)
+    return snake_case(value)
+
+
 def camel_case(value):
     return "".join(x.title() for x in snake_case(value).split("_"))
 
@@ -71,3 +95,245 @@ def header(self, text, level, raw=None):
 
 def docstring(text):
     return m2r2.convert(text.replace("\\n", "\\\\n"), renderer=CustomRenderer())[1:-1].replace("\\ ", " ")
+
+
+def _merge_imports(a, b):
+    """Merge second set of imports into first one."""
+    for k, v in b.items():
+        a[k] |= v
+    return a
+
+
+def format_parameters(kwargs, spec, version, replace_values=None):
+    parameters = ""
+    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"]
+    ):
+        parent = spec["requestBody"]["content"]["multipart/form-data"]["schema"]
+        for name, schema in parent["properties"].items():
+            parameters_spec[name] = {
+                "in": "form",
+                "schema": schema,
+                "name": name,
+                "description": schema.get("description"),
+                "required": name in parent.get("required", []),
+            }
+
+    parameters = ""
+    for p in parameters_spec.values():
+        k = p["name"]
+        if k not in kwargs:
+            continue
+
+        v = kwargs[k]
+        value, extra_imports = format_data_with_schema(
+            v["value"],
+            p["schema"],
+            replace_values=replace_values,
+            version=version,
+        )
+        imports = _merge_imports(imports, extra_imports)
+        parameters += f"{escape_reserved_keyword(safe_snake_case(k))}={value}, "
+
+    return parameters, imports
+
+
+def get_name_and_imports(schema, version=None, imports=None):
+    assert version is not None
+    imports = imports or defaultdict(set)
+
+    name = None
+    if hasattr(schema, "__reference__"):
+        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)
+
+    return name, imports
+
+
+@singledispatch
+def format_data_with_schema(
+    data,
+    schema,
+    replace_values=None,
+    default_name=None,
+    version=None,
+    imports=None,
+):
+    """Format data with schema."""
+    assert version is not None
+
+    name = None
+    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 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']}")
+
+    if replace_values and data in replace_values:
+        parameters = replace_values[data]
+        if schema.get("format") in ("int32", "int64"):
+            parameters = f"int({parameters})"
+    else:
+        if schema.get("nullable") and data is None:
+            parameters = repr(data)
+        else:
+
+            def format_datetime(x):
+                imports["datetime"].add("datetime")
+                d = dateutil.parser.isoparse(x)
+                result = repr(d)
+                if result.startswith("datetime."):
+                    result = result[len("datetime.") :]
+                if "tzutc" in result:
+                    imports["dateutil.tz"].add("tzutc")
+                return result
+
+            formatter = {
+                "double": lambda s: repr(float(s)),
+                "int32": lambda s: repr(int(s)),
+                "int64": lambda s: repr(int(s)),
+                "date": format_datetime,
+                "date-time": format_datetime,
+                "binary": lambda s: f'open("{s}", "rb")',
+                "email": repr,
+                None: repr,
+            }[schema.get("format")]
+
+            # TODO format date and datetime
+            parameters = formatter(data)
+
+    if name:
+        return f"{name}({parameters})", imports
+
+    return parameters, imports
+
+
+@format_data_with_schema.register(list)
+def format_data_with_schema_list(
+    data,
+    schema,
+    replace_values=None,
+    default_name=None,
+    version=None,
+    imports=None,
+):
+    """Format data with schema."""
+    assert version is not None
+    name, imports = get_name_and_imports(schema, version, imports)
+
+    parameters = ""
+    for d in data:
+        value, extra_imports = format_data_with_schema(
+            d,
+            schema["items"],
+            replace_values=replace_values,
+            default_name=name,
+            version=version,
+        )
+        parameters += f"{value}, "
+        imports = _merge_imports(imports, extra_imports)
+    parameters = f"[{parameters}]"
+
+    if name:
+        return f"{name}({parameters})", imports
+
+    return parameters, imports
+
+
+@format_data_with_schema.register(dict)
+def format_data_with_schema_dict(
+    data,
+    schema,
+    replace_values=None,
+    default_name=None,
+    version=None,
+    imports=None,
+):
+    """Format data with schema."""
+    assert version is not None
+    name, imports = get_name_and_imports(schema, version, imports)
+
+    parameters = ""
+    if "properties" in schema:
+        for k, v in data.items():
+            if k in schema["properties"]:
+                sub_schema = schema["properties"][k]
+            else:
+                sub_schema = schema["additionalProperties"]
+            value, extra_imports = format_data_with_schema(
+                v,
+                sub_schema,
+                replace_values=replace_values,
+                default_name=name + camel_case(k) if name else None,
+                version=version,
+            )
+            parameters += f"{escape_reserved_keyword(safe_snake_case(k))}={value}, "
+            imports = _merge_imports(imports, extra_imports)
+
+    if schema.get("additionalProperties") and not schema.get("properties"):
+        for k, v in data.items():
+            value, extra_imports = format_data_with_schema(
+                v,
+                schema["additionalProperties"],
+                replace_values=replace_values,
+                version=version,
+            )
+            parameters += f"{escape_reserved_keyword(k)}={value}, "
+            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")
+        ):
+            name = default_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}")
+
+    if "oneOf" in schema:
+        matched = 0
+        for sub_schema in schema["oneOf"]:
+            try:
+                formatted, extra_imports = format_data_with_schema(
+                    data,
+                    sub_schema,
+                    replace_values=replace_values,
+                    version=version,
+                )
+                if matched == 0:
+                    imports = _merge_imports(imports, extra_imports)
+                    # NOTE we do not support mixed schemas with oneOf
+                    # parameters += formatted
+                    parameters = formatted
+                    name = None
+                matched += 1
+            except (KeyError, ValueError) as e:
+                print(f"{e}")
+
+        if matched == 0:
+            raise ValueError(f"[{matched}] {data} is not valid for schema {name}")
+        elif matched > 1:
+            warnings.warn(f"[{matched}] {data} is not valid for schema {name}")
+
+    if name:
+        return f"{name}({parameters})", imports
+
+    return parameters, imports