From 0a16cfce5b383b2d67d066ae2de083bff72fc472 Mon Sep 17 00:00:00 2001
From: Roman Donchenko <roman@cvat.ai>
Date: Sun, 27 Nov 2022 13:40:19 +0300
Subject: [PATCH] SDK: add a high-level method to download task data chunks
 (#5356)

---
 cvat-sdk/cvat_sdk/core/proxies/tasks.py | 23 ++++++++++++++++++++++-
 tests/python/sdk/test_tasks.py          | 13 +++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/cvat-sdk/cvat_sdk/core/proxies/tasks.py b/cvat-sdk/cvat_sdk/core/proxies/tasks.py
index b5ae0c57..c0e89bb6 100644
--- a/cvat-sdk/cvat_sdk/core/proxies/tasks.py
+++ b/cvat-sdk/cvat_sdk/core/proxies/tasks.py
@@ -9,9 +9,10 @@ import json
 import mimetypes
 import os
 import os.path as osp
+import shutil
 from enum import Enum
 from time import sleep
-from typing import Any, Dict, List, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence
 
 from PIL import Image
 
@@ -32,6 +33,9 @@ from cvat_sdk.core.proxies.model_proxy import (
 from cvat_sdk.core.uploading import AnnotationUploader, DataUploader, Uploader
 from cvat_sdk.core.utils import filter_dict
 
+if TYPE_CHECKING:
+    from _typeshed import SupportsWrite
+
 
 class ResourceType(Enum):
     LOCAL = 0
@@ -153,6 +157,23 @@ class Task(
         (_, response) = self.api.retrieve_data(self.id, type="preview")
         return io.BytesIO(response.data)
 
+    def download_chunk(
+        self,
+        chunk_id: int,
+        output_file: SupportsWrite[bytes],
+        *,
+        quality: Optional[str] = None,
+    ) -> None:
+        params = {}
+        if quality:
+            params["quality"] = quality
+        (_, response) = self.api.retrieve_data(
+            self.id, number=chunk_id, **params, type="chunk", _parse_response=False
+        )
+
+        with response:
+            shutil.copyfileobj(response, output_file)
+
     def download_frames(
         self,
         frame_ids: Sequence[int],
diff --git a/tests/python/sdk/test_tasks.py b/tests/python/sdk/test_tasks.py
index e41ffde7..c4dcc136 100644
--- a/tests/python/sdk/test_tasks.py
+++ b/tests/python/sdk/test_tasks.py
@@ -4,6 +4,7 @@
 
 import io
 import os.path as osp
+import zipfile
 from logging import Logger
 from pathlib import Path
 from typing import Tuple
@@ -268,6 +269,18 @@ class TestTaskUsecases:
         assert osp.isfile(self.tmp_path / "frame-0.jpg")
         assert self.stdout.getvalue() == ""
 
+    @pytest.mark.parametrize("quality", ("compressed", "original"))
+    def test_can_download_chunk(self, fxt_new_task: Task, quality: str):
+        chunk_path = self.tmp_path / "chunk.zip"
+
+        with open(chunk_path, "wb") as chunk_file:
+            fxt_new_task.download_chunk(0, chunk_file, quality=quality)
+
+        with zipfile.ZipFile(chunk_path, "r") as chunk_zip:
+            assert chunk_zip.testzip() is None
+            assert len(chunk_zip.infolist()) == 1
+        assert self.stdout.getvalue() == ""
+
     def test_can_upload_annotations(self, fxt_new_task: Task, fxt_coco_file: Path):
         pbar_out = io.StringIO()
         pbar = make_pbar(file=pbar_out)