mapillary-tools 0.14.0a1__py3-none-any.whl → 0.14.0b1__py3-none-any.whl

mapillary_tools/upload_api_v4.py

@@ -6,15 +6,18 @@ import random
 import sys
 import typing as T
 import uuid
+from pathlib import Path
 
 if sys.version_info >= (3, 12):
     from typing import override
 else:
     from typing_extensions import override
 
+import tempfile
+
 import requests
 
-from .api_v4 import ClusterFileType, request_get, request_post, REQUESTS_TIMEOUT
+from .api_v4 import request_get, request_post, REQUESTS_TIMEOUT
 
 MAPILLARY_UPLOAD_ENDPOINT = os.getenv(
     "MAPILLARY_UPLOAD_ENDPOINT", "https://rupload.facebook.com/mapillary_public_uploads"
@@ -29,37 +32,23 @@ UPLOAD_REQUESTS_TIMEOUT = (30 * 60, 30 * 60)  # 30 minutes
 
 
 class UploadService:
+    """
+    Upload byte streams to the Upload Service.
+    """
+
     user_access_token: str
     session_key: str
-    cluster_filetype: ClusterFileType
-
-    MIME_BY_CLUSTER_TYPE: dict[ClusterFileType, str] = {
-        ClusterFileType.ZIP: "application/zip",
-        ClusterFileType.BLACKVUE: "video/mp4",
-        ClusterFileType.CAMM: "video/mp4",
-    }
 
-    def __init__(
-        self,
-        user_access_token: str,
-        session_key: str,
-        cluster_filetype: ClusterFileType,
-    ):
+    def __init__(self, user_access_token: str, session_key: str):
         self.user_access_token = user_access_token
         self.session_key = session_key
-        # Validate the input
-        self.cluster_filetype = cluster_filetype
 
     def fetch_offset(self) -> int:
         headers = {
             "Authorization": f"OAuth {self.user_access_token}",
         }
         url = f"{MAPILLARY_UPLOAD_ENDPOINT}/{self.session_key}"
-        resp = request_get(
-            url,
-            headers=headers,
-            timeout=REQUESTS_TIMEOUT,
-        )
+        resp = request_get(url, headers=headers, timeout=REQUESTS_TIMEOUT)
         resp.raise_for_status()
         data = resp.json()
         return data["offset"]
@@ -68,18 +57,53 @@ class UploadService:
     def chunkize_byte_stream(
         cls, stream: T.IO[bytes], chunk_size: int
     ) -> T.Generator[bytes, None, None]:
+        """
+        Chunkize a byte stream into chunks of the specified size.
+
+        >>> list(UploadService.chunkize_byte_stream(io.BytesIO(b"foo"), 1))
+        [b'f', b'o', b'o']
+
+        >>> list(UploadService.chunkize_byte_stream(io.BytesIO(b"foo"), 10))
+        [b'foo']
+        """
+
         if chunk_size <= 0:
             raise ValueError("Expect positive chunk size")
+
         while True:
             data = stream.read(chunk_size)
             if not data:
                 break
             yield data
 
+    @classmethod
     def shift_chunks(
-        self, chunks: T.Iterable[bytes], offset: int
+        cls, chunks: T.Iterable[bytes], offset: int
     ) -> T.Generator[bytes, None, None]:
-        assert offset >= 0, f"Expect non-negative offset but got {offset}"
+        """
+        Shift the chunks by the offset.
+
+        >>> list(UploadService.shift_chunks([b"foo", b"bar"], 0))
+        [b'foo', b'bar']
+
+        >>> list(UploadService.shift_chunks([b"foo", b"bar"], 1))
+        [b'oo', b'bar']
+
+        >>> list(UploadService.shift_chunks([b"foo", b"bar"], 3))
+        [b'bar']
+
+        >>> list(UploadService.shift_chunks([b"foo", b"bar"], 6))
+        []
+
+        >>> list(UploadService.shift_chunks([b"foo", b"bar"], 7))
+        []
+
+        >>> list(UploadService.shift_chunks([], 0))
+        []
+        """
+
+        if offset < 0:
+            raise ValueError(f"Expect non-negative offset but got {offset}")
 
         for chunk in chunks:
             if offset:
@@ -112,26 +136,20 @@ class UploadService:
         return self.upload_shifted_chunks(shifted_chunks, offset)
 
     def upload_shifted_chunks(
-        self,
-        shifted_chunks: T.Iterable[bytes],
-        offset: int,
+        self, shifted_chunks: T.Iterable[bytes], offset: int
     ) -> str:
         """
-        Upload the chunks that must already be shifted by the offset (e.g. fp.seek(begin_offset, io.SEEK_SET))
+        Upload the chunks that must already be shifted by the offset (e.g. fp.seek(offset, io.SEEK_SET))
         """
 
         headers = {
             "Authorization": f"OAuth {self.user_access_token}",
             "Offset": f"{offset}",
             "X-Entity-Name": self.session_key,
-            "X-Entity-Type": self.MIME_BY_CLUSTER_TYPE[self.cluster_filetype],
         }
         url = f"{MAPILLARY_UPLOAD_ENDPOINT}/{self.session_key}"
         resp = request_post(
-            url,
-            headers=headers,
-            data=shifted_chunks,
-            timeout=UPLOAD_REQUESTS_TIMEOUT,
+            url, headers=headers, data=shifted_chunks, timeout=UPLOAD_REQUESTS_TIMEOUT
         )
 
         resp.raise_for_status()
@@ -147,18 +165,35 @@ class UploadService:
 
 # A mock class for testing only
 class FakeUploadService(UploadService):
-    def __init__(self, *args, **kwargs):
+    """
+    A mock upload service that simulates the upload process for testing purposes.
+    It writes the uploaded data to a file in a temporary directory and generates a fake file handle.
+    """
+
+    FILE_HANDLE_DIR: str = "file_handles"
+
+    def __init__(
+        self,
+        upload_path: Path | None = None,
+        transient_error_ratio: float = 0.0,
+        *args,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
-        self._upload_path = os.getenv(
-            "MAPILLARY_UPLOAD_PATH", "mapillary_public_uploads"
-        )
-        self._error_ratio = 0.02
+        if upload_path is None:
+            upload_path = Path(tempfile.gettempdir()).joinpath(
+                "mapillary_public_uploads"
+            )
+        self._upload_path = upload_path
+        self._transient_error_ratio = transient_error_ratio
+
+    @property
+    def upload_path(self) -> Path:
+        return self._upload_path
 
     @override
     def upload_shifted_chunks(
-        self,
-        shifted_chunks: T.Iterable[bytes],
-        offset: int,
+        self, shifted_chunks: T.Iterable[bytes], offset: int
     ) -> str:
         expected_offset = self.fetch_offset()
         if offset != expected_offset:
@@ -167,28 +202,36 @@ class FakeUploadService(UploadService):
             )
 
         os.makedirs(self._upload_path, exist_ok=True)
-        filename = os.path.join(self._upload_path, self.session_key)
-        with open(filename, "ab") as fp:
+        filename = self._upload_path.joinpath(self.session_key)
+        with filename.open("ab") as fp:
             for chunk in shifted_chunks:
-                if random.random() <= self._error_ratio:
+                if random.random() <= self._transient_error_ratio:
                     raise requests.ConnectionError(
-                        f"TEST ONLY: Failed to upload with error ratio {self._error_ratio}"
+                        f"TEST ONLY: Failed to upload with error ratio {self._transient_error_ratio}"
                     )
                 fp.write(chunk)
-                if random.random() <= self._error_ratio:
+                if random.random() <= self._transient_error_ratio:
                     raise requests.ConnectionError(
-                        f"TEST ONLY: Partially uploaded with error ratio {self._error_ratio}"
+                        f"TEST ONLY: Partially uploaded with error ratio {self._transient_error_ratio}"
                     )
-        return uuid.uuid4().hex
+
+        file_handle_dir = self._upload_path.joinpath(self.FILE_HANDLE_DIR)
+        file_handle_path = file_handle_dir.joinpath(self.session_key)
+        if not file_handle_path.exists():
+            os.makedirs(file_handle_dir, exist_ok=True)
+            random_file_handle = uuid.uuid4().hex
+            file_handle_path.write_text(random_file_handle)
+
+        return file_handle_path.read_text()
 
     @override
     def fetch_offset(self) -> int:
-        if random.random() <= self._error_ratio:
+        if random.random() <= self._transient_error_ratio:
             raise requests.ConnectionError(
-                f"TEST ONLY: Partially uploaded with error ratio {self._error_ratio}"
+                f"TEST ONLY: Partially uploaded with error ratio {self._transient_error_ratio}"
             )
-        filename = os.path.join(self._upload_path, self.session_key)
-        if not os.path.exists(filename):
+        filename = self._upload_path.joinpath(self.session_key)
+        if not filename.exists():
             return 0
         with open(filename, "rb") as fp:
             fp.seek(0, io.SEEK_END)