Compute content length

Author: JordonPhillips

packages/smithy-core/src/smithy_core/traits.py

@@ -154,6 +154,12 @@ def __post_init__(self):
         assert self.document_value is None
 
 
+@dataclass(init=False, frozen=True)
+class RequiresLengthTrait(Trait, id=ShapeID("smithy.api#requiresLength")):
+    def __post_init__(self):
+        assert self.document_value is None
+
+
 @dataclass(init=False, frozen=True)
 class UnitTypeTrait(Trait, id=ShapeID("smithy.api#UnitTypeTrait")):
     def __post_init__(self):

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

@@ -9,6 +9,7 @@
 
 from smithy_core import URI
 from smithy_core.codecs import Codec
+from smithy_core.exceptions import SerializationError
 from smithy_core.schemas import Schema
 from smithy_core.serializers import (
     InterceptingSerializer,
@@ -24,12 +25,13 @@
     HTTPQueryTrait,
     HTTPTrait,
     MediaTypeTrait,
+    RequiresLengthTrait,
     TimestampFormatTrait,
 )
 from smithy_core.types import PathPattern, TimestampFormat
 from smithy_core.utils import serialize_float
 
-from . import tuples_to_fields
+from . import Field, tuples_to_fields
 from .aio import HTTPRequest as _HTTPRequest
 from .aio import HTTPResponse as _HTTPResponse
 from .aio.interfaces import HTTPRequest, HTTPResponse
@@ -43,6 +45,7 @@
 __all__ = ["HTTPRequestSerializer", "HTTPResponseSerializer"]
 
 
+# TODO: refactor this to share code with response serializer
 class HTTPRequestSerializer(SpecificShapeSerializer):
     """Binds a serializable shape to an HTTP request.
 
@@ -82,8 +85,12 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
             host_prefix = self._endpoint_trait.host_prefix
 
         content_type = self._payload_codec.media_type
+        content_length: int | None = None
+        content_length_required = False
+
         binding_matcher = RequestBindingMatcher(schema)
         if (payload_member := binding_matcher.payload_member) is not None:
+            content_length_required = RequiresLengthTrait in payload_member
             if payload_member.shape_type in (
                 ShapeType.BLOB,
                 ShapeType.STRING,
@@ -105,6 +112,10 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
                 )
                 yield binding_serializer
                 payload = payload_serializer.payload
+                try:
+                    content_length = len(payload)
+                except TypeError:
+                    pass
             else:
                 if (media_type := payload_member.get_trait(MediaTypeTrait)) is not None:
                     content_type = media_type.value
@@ -117,6 +128,7 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
                     binding_matcher,
                 )
                 yield binding_serializer
+                content_length = payload.tell()
         else:
             payload = BytesIO()
             payload_serializer = self._payload_codec.create_serializer(payload)
@@ -131,8 +143,10 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
                         binding_matcher,
                     )
                     yield binding_serializer
+                content_length = payload.tell()
             else:
                 content_type = None
+                content_length = 0
                 binding_serializer = HTTPRequestBindingSerializer(
                     payload_serializer,
                     self._http_trait.path,
@@ -141,15 +155,24 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
                 )
                 yield binding_serializer
 
-        if (
-            seek := getattr(payload, "seek", None)
-        ) is not None and not iscoroutinefunction(seek):
-            seek(0)
-
+        self._seek(payload, 0)
         headers = binding_serializer.header_serializer.headers
         if content_type is not None:
             headers.append(("content-type", content_type))
 
+        if content_length is not None:
+            headers.append(("content-length", str(content_length)))
+
+        fields = tuples_to_fields(headers)
+        if content_length_required and "content-length" not in fields:
+            content_length = self._compute_content_length(payload)
+            if content_length is None:
+                raise SerializationError(
+                    "This operation requires the the content length of the input "
+                    "stream, but it was not provided and was unable to be computed."
+                )
+            fields.set_field(Field(name="content-length", values=[str(content_length)]))
+
         self.result = _HTTPRequest(
             method=self._http_trait.method,
             destination=URI(
@@ -160,10 +183,22 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
                     prefix=self._http_trait.query or "",
                 ),
             ),
-            fields=tuples_to_fields(headers),
+            fields=fields,
             body=payload,
         )
 
+    def _seek(self, payload: Any, pos: int, whence: int = 0) -> None:
+        if (
+            seek := getattr(payload, "seek", None)
+        ) is not None and not iscoroutinefunction(seek):
+            seek(pos, whence)
+
+    def _compute_content_length(self, payload: Any) -> int | None:
+        content_length = self._seek(payload, 0, 2)
+        if content_length is not None:
+            self._seek(payload, 0, 0)
+            return content_length
+
 
 class HTTPRequestBindingSerializer(InterceptingSerializer):
     """Delegates HTTP request bindings to binding-location-specific serializers."""
@@ -235,8 +270,12 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
         binding_serializer: HTTPResponseBindingSerializer
 
         content_type: str | None = self._payload_codec.media_type
+        content_length: int | None = None
+        content_length_required = False
+
         binding_matcher = ResponseBindingMatcher(schema)
         if (payload_member := binding_matcher.payload_member) is not None:
+            content_length_required = RequiresLengthTrait in payload_member
             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
@@ -250,6 +289,10 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
                 )
                 yield binding_serializer
                 payload = payload_serializer.payload
+                try:
+                    content_length = len(payload)
+                except TypeError:
+                    pass
             else:
                 if (media_type := payload_member.get_trait(MediaTypeTrait)) is not None:
                     content_type = media_type.value
@@ -259,6 +302,7 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
                     payload_serializer, binding_matcher
                 )
                 yield binding_serializer
+                content_length = payload.tell()
         else:
             payload = BytesIO()
             payload_serializer = self._payload_codec.create_serializer(payload)
@@ -270,23 +314,34 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
                         body_serializer, binding_matcher
                     )
                     yield binding_serializer
+                content_length = payload.tell()
             else:
                 content_type = None
+                content_length = 0
                 binding_serializer = HTTPResponseBindingSerializer(
                     payload_serializer,
                     binding_matcher,
                 )
                 yield binding_serializer
 
-        if (
-            seek := getattr(payload, "seek", None)
-        ) is not None and not iscoroutinefunction(seek):
-            seek(0)
-
+        self._seek(payload, 0)
         headers = binding_serializer.header_serializer.headers
         if content_type is not None:
             headers.append(("content-type", content_type))
 
+        if content_length is not None:
+            headers.append(("content-length", str(content_length)))
+
+        fields = tuples_to_fields(headers)
+        if content_length_required and "content-length" not in fields:
+            content_length = self._compute_content_length(payload)
+            if content_length is None:
+                raise SerializationError(
+                    "This operation requires the the content length of the input "
+                    "stream, but it was not provided and was unable to be computed."
+                )
+            fields.set_field(Field(name="content-length", values=[str(content_length)]))
+
         status = binding_serializer.response_code_serializer.response_code
         if status is None:
             if binding_matcher.response_status > 0:
@@ -300,6 +355,18 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
             status=status,
         )
 
+    def _seek(self, payload: Any, pos: int, whence: int = 0) -> int | None:
+        if (
+            seek := getattr(payload, "seek", None)
+        ) is not None and not iscoroutinefunction(seek):
+            return seek(pos, whence)
+
+    def _compute_content_length(self, payload: Any) -> int | None:
+        content_length = self._seek(payload, 0, 2)
+        if content_length is not None:
+            self._seek(payload, 0, 0)
+            return content_length
+
 
 class HTTPResponseBindingSerializer(InterceptingSerializer):
     """Delegates HTTP response bindings to binding-location-specific serializers."""