Only write content type when necessary

Author: JordonPhillips

packages/smithy-http/src/smithy_http/serializers.py

@@ -82,15 +82,20 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
             host_prefix = self._endpoint_trait.host_prefix
 
         content_type = self._payload_codec.media_type
-
         binding_matcher = RequestBindingMatcher(schema)
         if (payload_member := binding_matcher.payload_member) is not None:
-            if payload_member.shape_type in (ShapeType.BLOB, ShapeType.STRING):
-                content_type = (
-                    "application/octet-stream"
-                    if payload_member.shape_type is ShapeType.BLOB
-                    else "text/plain"
-                )
+            if payload_member.shape_type in (
+                ShapeType.BLOB,
+                ShapeType.STRING,
+                ShapeType.ENUM,
+            ):
+                if (media_type := payload_member.get_trait(MediaTypeTrait)) is not None:
+                    content_type = media_type.value
+                elif payload_member.shape_type is ShapeType.BLOB:
+                    content_type = "application/octet-stream"
+                else:
+                    content_type = "text/plain"
+
                 payload_serializer = RawPayloadSerializer()
                 binding_serializer = HTTPRequestBindingSerializer(
                     payload_serializer,
@@ -113,11 +118,11 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
                 )
                 yield binding_serializer
         else:
-            if binding_matcher.event_stream_member is not None:
-                content_type = "application/vnd.amazon.eventstream"
             payload = BytesIO()
             payload_serializer = self._payload_codec.create_serializer(payload)
             if binding_matcher.should_write_body(self._omit_empty_payload):
+                if binding_matcher.event_stream_member is not None:
+                    content_type = "application/vnd.amazon.eventstream"
                 with payload_serializer.begin_struct(schema) as body_serializer:
                     binding_serializer = HTTPRequestBindingSerializer(
                         body_serializer,
@@ -127,6 +132,7 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
                     )
                     yield binding_serializer
             else:
+                content_type = None
                 binding_serializer = HTTPRequestBindingSerializer(
                     payload_serializer,
                     self._http_trait.path,
@@ -140,9 +146,9 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
         ) is not None and not iscoroutinefunction(seek):
             seek(0)
 
-        # TODO: conditional on empty-ness and based on the protocol
         headers = binding_serializer.header_serializer.headers
-        headers.append(("content-type", content_type))
+        if content_type is not None:
+            headers.append(("content-type", content_type))
 
         self.result = _HTTPRequest(
             method=self._http_trait.method,
@@ -228,16 +234,25 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
         payload: Any
         binding_serializer: HTTPResponseBindingSerializer
 
+        content_type: str | None = self._payload_codec.media_type
         binding_matcher = ResponseBindingMatcher(schema)
         if (payload_member := binding_matcher.payload_member) is not None:
             if payload_member.shape_type in (ShapeType.BLOB, ShapeType.STRING):
+                if (media_type := payload_member.get_trait(MediaTypeTrait)) is not None:
+                    content_type = media_type.value
+                elif payload_member.shape_type is ShapeType.BLOB:
+                    content_type = "application/octet-stream"
+                else:
+                    content_type = "text/plain"
                 payload_serializer = RawPayloadSerializer()
                 binding_serializer = HTTPResponseBindingSerializer(
                     payload_serializer, binding_matcher
                 )
                 yield binding_serializer
                 payload = payload_serializer.payload
             else:
+                if (media_type := payload_member.get_trait(MediaTypeTrait)) is not None:
+                    content_type = media_type.value
                 payload = BytesIO()
                 payload_serializer = self._payload_codec.create_serializer(payload)
                 binding_serializer = HTTPResponseBindingSerializer(
@@ -248,12 +263,15 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
             payload = BytesIO()
             payload_serializer = self._payload_codec.create_serializer(payload)
             if binding_matcher.should_write_body(self._omit_empty_payload):
+                if binding_matcher.event_stream_member is not None:
+                    content_type = "application/vnd.amazon.eventstream"
                 with payload_serializer.begin_struct(schema) as body_serializer:
                     binding_serializer = HTTPResponseBindingSerializer(
                         body_serializer, binding_matcher
                     )
                     yield binding_serializer
             else:
+                content_type = None
                 binding_serializer = HTTPResponseBindingSerializer(
                     payload_serializer,
                     binding_matcher,
@@ -265,6 +283,10 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
         ) is not None and not iscoroutinefunction(seek):
             seek(0)
 
+        headers = binding_serializer.header_serializer.headers
+        if content_type is not None:
+            headers.append(("content-type", content_type))
+
         status = binding_serializer.response_code_serializer.response_code
         if status is None:
             if binding_matcher.response_status > 0:

packages/smithy-http/tests/unit/test_serializers.py

@@ -40,7 +40,7 @@
     TimestampFormatTrait,
     Trait,
 )
-from smithy_http import Field, Fields, tuples_to_fields
+from smithy_http import Fields, tuples_to_fields
 from smithy_http.aio import HTTPResponse as _HTTPResponse
 from smithy_http.deserializers import HTTPResponseDeserializer
 from smithy_http.serializers import HTTPRequestSerializer, HTTPResponseSerializer
@@ -1552,7 +1552,9 @@ def host_cases() -> list[HTTPMessageTestCase]:
         HTTPMessageTestCase(
             HostLabel("foo"),
             HTTPMessage(
-                destination=URI(host="foo.", path="/"), body=BytesIO(b'{"label":"foo"}')
+                destination=URI(host="foo.", path="/"),
+                body=BytesIO(b'{"label":"foo"}'),
+                fields=tuples_to_fields([("content-type", "application/json")]),
             ),
             endpoint_trait=EndpointTrait({"hostPrefix": "{label}."}),
         ),
@@ -1569,7 +1571,9 @@ def payload_cases() -> list[HTTPMessageTestCase]:
         HTTPMessageTestCase(
             HTTPImplicitPayload(header="foo", payload_member="bar"),
             HTTPMessage(
-                fields=tuples_to_fields([("header", "foo")]),
+                fields=tuples_to_fields(
+                    [("header", "foo"), ("content-type", "application/json")]
+                ),
                 body=BytesIO(b'{"payload_member":"bar"}'),
             ),
         ),
@@ -1588,7 +1592,10 @@ def payload_cases() -> list[HTTPMessageTestCase]:
         ),
         HTTPMessageTestCase(
             HTTPStructuredPayload(payload=HTTPStringPayload(payload="foo")),
-            HTTPMessage(body=BytesIO(b'{"payload":"foo"}')),
+            HTTPMessage(
+                body=BytesIO(b'{"payload":"foo"}'),
+                fields=tuples_to_fields([("content-type", "application/json")]),
+            ),
         ),
     ]
 
@@ -1615,8 +1622,6 @@ def async_streaming_payload_cases() -> list[HTTPMessageTestCase]:
     + async_streaming_payload_cases()
 )
 
-CONTENT_TYPE_FIELD = Field(name="content-type", values=["application/json"])
-
 
 @pytest.mark.parametrize("case", REQUEST_SER_CASES)
 async def test_serialize_http_request(case: HTTPMessageTestCase) -> None:
@@ -1636,10 +1641,6 @@ async def test_serialize_http_request(case: HTTPMessageTestCase) -> None:
     actual_query = actual.destination.query or ""
     expected_query = case.request.destination.query or ""
     assert actual_query == expected_query
-    # set the content-type field here, otherwise cases would have to duplicate it everywhere,
-    # but if the field is already set in the case, don't override it
-    if expected.fields.get(CONTENT_TYPE_FIELD.name) is None:
-        expected.fields.set_field(CONTENT_TYPE_FIELD)
     assert actual.fields == expected.fields
 
     if case.request.body:
@@ -1678,9 +1679,6 @@ async def test_serialize_http_response(case: HTTPMessageTestCase) -> None:
     expected = case.request
 
     assert actual is not None
-    # Remove content-type from expected, we're re-using the request cases for brevity
-    if expected.fields.get(CONTENT_TYPE_FIELD.name) is not None:
-        del expected.fields[CONTENT_TYPE_FIELD.name]
     assert actual.fields == expected.fields
     assert actual.status == expected.status