[PJL-10221] handle some edge cases in OpenAPI v3.1 swagger generation (…

…#319) * [PJL-10221] update some edge cases * updates * update tests * typo * fix type * typo
plangrid · Oct 10, 2024 · c9edc64 · c9edc64
1 parent e438006
commit c9edc64
Show file tree

Hide file tree

Showing 2 changed files with 103 additions and 0 deletions.
diff --git a/flask_rebar/swagger_generation/marshmallow_to_swagger.py b/flask_rebar/swagger_generation/marshmallow_to_swagger.py
@@ -37,6 +37,7 @@
 from marshmallow import Schema
 from marshmallow.validate import Range
 from marshmallow.validate import OneOf
+from marshmallow.validate import ContainsOnly
 from marshmallow.validate import Length
 from marshmallow.validate import Validator
 
@@ -423,6 +424,15 @@ class DictConverter(FieldConverter[m.fields.Dict]):
     def get_type(self, obj: m.fields.Dict, context: _Context) -> Union[str, List[str]]:
         return self.null_type_determination(obj, context, sw.object_)
 
+    @sets_swagger_attr(sw.additional_properties)
+    def get_additional_properties(
+        self, obj: m.fields.Dict, context: _Context
+    ) -> Union[Type[UNSET], m.fields.Dict]:
+        if obj.value_field:
+            return context.convert(obj.value_field, context)
+        else:
+            return UNSET
+
 
 class IntegerConverter(FieldConverter[m.fields.Integer]):
     MARSHMALLOW_TYPE = m.fields.Integer
@@ -604,6 +614,14 @@ def get_maximum(
             return UNSET
 
 
+class ContainsOnlyConverter(ValidatorConverter):
+    MARSHMALLOW_TYPE = ContainsOnly
+
+    @sets_swagger_attr(sw.items)
+    def get_items(self, obj: ContainsOnly, context: _Context) -> Dict[str, Any]:
+        return {"type": context.memo["items"]["type"], "enum": obj.choices}
+
+
 class OneOfConverter(ValidatorConverter):
     MARSHMALLOW_TYPE = OneOf
 
@@ -794,6 +812,7 @@ def _common_converters() -> List[MarshmallowConverter]:
     """Instantiates the converters we use in ALL of the registries below"""
     converters: List[MarshmallowConverter] = [
         BooleanConverter(),
+        ContainsOnlyConverter(),
         DateConverter(),
         DateTimeConverter(),
         FunctionConverter(),

diff --git a/tests/swagger_generation/test_marshmallow_to_swagger.py b/tests/swagger_generation/test_marshmallow_to_swagger.py
@@ -195,6 +195,90 @@ class Foo(m.Schema):
                 },
             )
 
+    def test_list_enum_openapi_v3(self):
+        for field, result in [
+            (
+                QueryParamList(m.fields.Integer(), validate=v.ContainsOnly([1, 2, 3])),
+                {
+                    "type": "array",
+                    "items": {"type": "integer", "enum": [1, 2, 3]},
+                    "explode": True,
+                },
+            ),
+            (
+                CommaSeparatedList(
+                    m.fields.String(), validate=v.ContainsOnly(["a", "b", "c"])
+                ),
+                {
+                    "type": "array",
+                    "items": {"type": "string", "enum": ["a", "b", "c"]},
+                    "style": "form",
+                    "explode": False,
+                },
+            ),
+        ]:
+
+            class Foo(m.Schema):
+                a = field
+
+            schema = Foo()
+            json_schema = self.registry.convert(schema, openapi_version=3)
+
+            self.assertEqual(
+                json_schema,
+                {
+                    "additionalProperties": False,
+                    "type": "object",
+                    "title": "Foo",
+                    "properties": {"a": result},
+                },
+            )
+
+    def test_custom_dicts_openapi_v3(self):
+        for field, result in [
+            (
+                m.fields.Dict(),
+                {
+                    "type": "object",
+                },
+            ),
+            (
+                m.fields.Dict(
+                    values=m.fields.Integer(),
+                ),
+                {
+                    "type": "object",
+                    "additionalProperties": {"type": "integer"},
+                },
+            ),
+            (
+                m.fields.Dict(
+                    keys=m.fields.String(),
+                    values=m.fields.String(),
+                ),
+                {
+                    "type": "object",
+                    "additionalProperties": {"type": "string"},
+                },
+            ),
+        ]:
+
+            class Foo(m.Schema):
+                a = field
+
+            schema = Foo()
+            json_schema = self.registry.convert(schema, openapi_version=3)
+
+            self.assertEqual(
+                json_schema,
+                {
+                    "additionalProperties": False,
+                    "type": "object",
+                    "title": "Foo",
+                    "properties": {"a": result},
+                },
+            )
+
     def test_data_key(self):
         registry = ConverterRegistry()
         registry.register_types(ALL_CONVERTERS)