Add support for creation-defer-length extension (#96)

* add declare length option

* implement upload_length_deferred

* fix declare_length compatibility with upload_length_deferred

* clarify upload_length_deferred description

Co-authored-by: Marius Kleidl <[email protected]>

* remove declare_length

---------

Co-authored-by: Dmitriy Brusenin <[email protected]>
Co-authored-by: Marius Kleidl <[email protected]>

Author: 3 people

tests/test_uploader.py

@@ -4,6 +4,7 @@
 from unittest import mock
 
 import responses
+from responses import matchers
 from parametrize import parametrize
 import pytest
 
@@ -209,3 +210,42 @@ def test_upload_checksum(self, request_mock):
         self.uploader.upload_checksum = True
         self.uploader.upload()
         self.assertEqual(self.uploader.offset, self.uploader.get_file_size())
+    
+    @parametrize("chunk_size", [1, 2, 3, 4, 5, 6])
+    @responses.activate
+    def test_upload_length_deferred(self, chunk_size: int):
+        upload_url = f"{self.client.url}test_upload_length_deferred"
+        
+        responses.head(
+            upload_url,
+            adding_headers={"upload-offset": "0", "Upload-Defer-Length": "1"},
+        )
+        uploader = self.client.uploader(
+            file_stream=io.BytesIO(b"hello"),
+            url=upload_url,
+            chunk_size=chunk_size,
+            upload_length_deferred=True,
+        )
+        self.assertTrue(uploader.upload_length_deferred)
+        self.assertTrue(uploader.stop_at is None)
+        
+        offset = 0
+        while not (offset + chunk_size > 5):
+            next_offset = min(offset + chunk_size, 5)
+            responses.patch(
+                upload_url,
+                adding_headers={"upload-offset": str(next_offset)},
+                match=[matchers.header_matcher({"upload-offset": str(offset)})],
+            )
+            offset = next_offset
+        last_req_headers = {"upload-offset": str(offset)}
+        last_req_headers["upload-length"] = "5"
+        responses.patch(
+            upload_url,
+            adding_headers={"upload-offset": "5"},
+            match=[matchers.header_matcher(last_req_headers)],
+        )
+        
+        uploader.upload()
+        self.assertEqual(uploader.offset, 5)
+        self.assertEqual(uploader.stop_at, 5)

tusclient/request.py

@@ -41,9 +41,10 @@ class BaseTusRequest:
 
     def __init__(self, uploader):
         self._url = uploader.url
-        self.response_headers = {}
         self.status_code = None
+        self.response_headers = {}
         self.response_content = None
+        self.stream_eof = False
         self.verify_tls_cert = bool(uploader.verify_tls_cert)
         self.file = uploader.get_file_stream()
         self.file.seek(uploader.offset)
@@ -53,6 +54,8 @@ def __init__(self, uploader):
             "upload-offset": str(uploader.offset),
             "Content-Type": "application/offset+octet-stream",
         }
+        self._offset = uploader.offset
+        self._upload_length_deferred = uploader.upload_length_deferred
         self._request_headers.update(uploader.get_headers())
         self._content_length = uploader.get_request_length()
         self._upload_checksum = uploader.upload_checksum
@@ -80,18 +83,23 @@ def perform(self):
         """
         try:
             chunk = self.file.read(self._content_length)
+            stream_eof = len(chunk) < self._content_length
             self.add_checksum(chunk)
+            headers = self._request_headers
+            if stream_eof and self._upload_length_deferred:
+                headers["upload-length"] = str(self._offset + len(chunk))
             resp = requests.patch(
                 self._url,
                 data=chunk,
-                headers=self._request_headers,
+                headers=headers,
                 verify=self.verify_tls_cert,
                 stream=True,
                 cert=self.client_cert
             )
             self.status_code = resp.status_code
             self.response_content = resp.content
             self.response_headers = {k.lower(): v for k, v in resp.headers.items()}
+            self.stream_eof = stream_eof
         except requests.exceptions.RequestException as error:
             raise TusUploadFailed(error)
 

tusclient/uploader/baseuploader.py

@@ -73,6 +73,10 @@ class BaseUploader:
         - upload_checksum (bool):
             Whether or not to supply the Upload-Checksum header along with each
             chunk. Defaults to False.
+        - upload_length_deferred (bool):
+            Whether or not to declare the upload length when finished reading the file stream instead of when the upload is started. This is useful
+            when uploading from a streaming resource, where the total file size isn't available when the upload is created
+            but only becomes known when the stream finishes. The server must support the `creation-defer-length` extension.
 
     :Constructor Args:
         - file_path (str)
@@ -89,6 +93,7 @@ class BaseUploader:
         - url_storage (Optinal [<tusclient.storage.interface.Storage>])
         - fingerprinter (Optional [<tusclient.fingerprint.interface.Fingerprint>])
         - upload_checksum (Optional[bool])
+        - upload_length_deferred (Optional[bool])
     """
 
     DEFAULT_HEADERS = {"Tus-Resumable": "1.0.0"}
@@ -114,6 +119,7 @@ def __init__(
         url_storage: Optional[Storage] = None,
         fingerprinter: Optional[interface.Fingerprint] = None,
         upload_checksum=False,
+        upload_length_deferred=False,
     ):
         if file_path is None and file_stream is None:
             raise ValueError("Either 'file_path' or 'file_stream' cannot be None.")
@@ -129,7 +135,8 @@ def __init__(
         self.verify_tls_cert = verify_tls_cert
         self.file_path = file_path
         self.file_stream = file_stream
-        self.stop_at = self.get_file_size()
+        self.file_size = self.get_file_size() if not upload_length_deferred else None
+        self.stop_at = self.file_size
         self.client = client
         self.metadata = metadata or {}
         self.metadata_encoding = metadata_encoding
@@ -145,6 +152,7 @@ def __init__(
         self._retried = 0
         self.retry_delay = retry_delay
         self.upload_checksum = upload_checksum
+        self.upload_length_deferred = upload_length_deferred
         (
             self.__checksum_algorithm_name,
             self.__checksum_algorithm,
@@ -161,7 +169,10 @@ def get_headers(self):
     def get_url_creation_headers(self):
         """Return headers required to create upload url"""
         headers = self.get_headers()
-        headers["upload-length"] = str(self.get_file_size())
+        if self.upload_length_deferred:
+            headers['upload-defer-length'] = '1'
+        else:
+            headers["upload-length"] = str(self.file_size)
         headers["upload-metadata"] = ",".join(self.encode_metadata())
         return headers
 
@@ -252,8 +263,9 @@ def get_request_length(self):
         """
         Return length of next chunk upload.
         """
-        remainder = self.stop_at - self.offset
-        return self.chunk_size if remainder > self.chunk_size else remainder
+        if self.stop_at is None:
+            return self.chunk_size
+        return min(self.chunk_size, self.stop_at - self.offset)
 
     def get_file_stream(self):
         """

tusclient/uploader/uploader.py

@@ -33,7 +33,7 @@ def upload(self, stop_at: Optional[int] = None):
                 Determines at what offset value the upload should stop. If not specified this
                 defaults to the file size.
         """
-        self.stop_at = stop_at or self.get_file_size()
+        self.stop_at = stop_at or self.file_size
 
         if not self.url:
             # Ensure the POST request is performed even for empty files.
@@ -42,7 +42,7 @@ def upload(self, stop_at: Optional[int] = None):
             self.set_url(self.create_url())
             self.offset = 0
 
-        while self.offset < self.stop_at:
+        while self.stop_at is None or (self.offset < self.stop_at):
             self.upload_chunk()
 
     def upload_chunk(self):
@@ -59,6 +59,8 @@ def upload_chunk(self):
 
         self._do_request()
         self.offset = int(self.request.response_headers.get("upload-offset"))
+        if self.upload_length_deferred and self.request.stream_eof:
+            self.stop_at = self.offset
 
     @catch_requests_error
     def create_url(self):
@@ -120,13 +122,13 @@ async def upload(self, stop_at: Optional[int] = None):
                 Determines at what offset value the upload should stop. If not specified this
                 defaults to the file size.
         """
-        self.stop_at = stop_at or self.get_file_size()
+        self.stop_at = stop_at or self.file_size
 
         if not self.url:
             self.set_url(await self.create_url())
             self.offset = 0
 
-        while self.offset < self.stop_at:
+        while self.stop_at is None or (self.offset < self.stop_at):
             await self.upload_chunk()
 
     async def upload_chunk(self):
@@ -143,6 +145,8 @@ async def upload_chunk(self):
 
         await self._do_request()
         self.offset = int(self.request.response_headers.get("upload-offset"))
+        if self.upload_length_deferred and self.request.stream_eof:
+            self.stop_at = self.offset
 
     async def create_url(self):
         """