feat(downloads): allow streaming downloads access to response iterator

TCatshoek · TCatshoek · commit 4f9807f2fc13 · 2022-04-09T14:37:58.000+02:00
Allow access to the underlying response iterator when downloading in
streaming mode by specifying action="iterator"

Update type annotations to support this change
diff --git a/gitlab/mixins.py b/gitlab/mixins.py
@@ -20,7 +20,9 @@
     Any,
     Callable,
     Dict,
+    Iterator,
     List,
+    Literal,
     Optional,
     Tuple,
     Type,
@@ -657,10 +659,10 @@ class DownloadMixin(_RestObjectBase):
     def download(
         self,
         streamed: bool = False,
-        action: Optional[Callable] = None,
+        action: Optional[Union[Callable, Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Download the archive of a resource export.
 
         Args:
diff --git a/gitlab/utils.py b/gitlab/utils.py
@@ -19,7 +19,7 @@
 import traceback
 import urllib.parse
 import warnings
-from typing import Any, Callable, Dict, Optional, Type, Union
+from typing import Any, Callable, Dict, Iterator, Literal, Optional, Type, Union
 
 import requests
 
@@ -32,15 +32,18 @@ def __call__(self, chunk: Any) -> None:
 def response_content(
     response: requests.Response,
     streamed: bool,
-    action: Optional[Callable],
+    action: Optional[Union[Callable, Literal["iterator"]]],
     chunk_size: int,
-) -> Optional[bytes]:
+) -> Optional[Union[bytes, Iterator[Any]]]:
     if streamed is False:
         return response.content
 
     if action is None:
         action = _StdoutStream()
 
+    if action == "iterator":
+        return response.iter_content(chunk_size=chunk_size)
+
     for chunk in response.iter_content(chunk_size=chunk_size):
         if chunk:
             action(chunk)
diff --git a/gitlab/v4/cli.py b/gitlab/v4/cli.py
@@ -19,7 +19,7 @@
 import argparse
 import operator
 import sys
-from typing import Any, Dict, List, Optional, Type, TYPE_CHECKING, Union
+from typing import Any, Dict, Iterator, List, Optional, Type, TYPE_CHECKING, Union
 
 import gitlab
 import gitlab.base
@@ -123,6 +123,7 @@ def do_project_export_download(self) -> None:
             data = export_status.download()
             if TYPE_CHECKING:
                 assert data is not None
+                assert not isinstance(data, Iterator)
             sys.stdout.buffer.write(data)
 
         except Exception as e:
diff --git a/gitlab/v4/objects/artifacts.py b/gitlab/v4/objects/artifacts.py
@@ -2,7 +2,7 @@
 GitLab API:
 https://docs.gitlab.com/ee/api/job_artifacts.html
 """
-from typing import Any, Callable, Optional, TYPE_CHECKING
+from typing import Any, Callable, Iterator, Literal, Optional, TYPE_CHECKING, Union
 
 import requests
 
@@ -32,7 +32,7 @@ def __call__(
         self,
         *args: Any,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         utils.warn(
             message=(
                 "The project.artifacts() method is deprecated and will be removed in a "
@@ -71,10 +71,10 @@ def download(
         ref_name: str,
         job: str,
         streamed: bool = False,
-        action: Optional[Callable] = None,
+        action: Optional[Union[Callable, Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Get the job artifacts archive from a specific tag or branch.
 
         Args:
@@ -86,7 +86,8 @@ def download(
                 `chunk_size` and each chunk is passed to `action` for
                 treatment
             action: Callable responsible of dealing with chunk of
-                data
+                data. May also be the string "iterator" to directly return
+                the response iterator
             chunk_size: Size of each chunk
             **kwargs: Extra options to send to the server (e.g. sudo)
 
@@ -115,10 +116,10 @@ def raw(
         artifact_path: str,
         job: str,
         streamed: bool = False,
-        action: Optional[Callable] = None,
+        action: Optional[Union[Callable, Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Download a single artifact file from a specific tag or branch from
         within the job's artifacts archive.
 
diff --git a/gitlab/v4/objects/files.py b/gitlab/v4/objects/files.py
@@ -1,5 +1,16 @@
 import base64
-from typing import Any, Callable, cast, Dict, List, Optional, TYPE_CHECKING
+from typing import (
+    Any,
+    Callable,
+    cast,
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    TYPE_CHECKING,
+    Union,
+)
 
 import requests
 
@@ -217,10 +228,10 @@ def raw(
         file_path: str,
         ref: str,
         streamed: bool = False,
-        action: Optional[Callable[..., Any]] = None,
+        action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Return the content of a file for a commit.
 
         Args:
diff --git a/gitlab/v4/objects/jobs.py b/gitlab/v4/objects/jobs.py
@@ -1,4 +1,14 @@
-from typing import Any, Callable, cast, Dict, Optional, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Callable,
+    cast,
+    Dict,
+    Iterator,
+    Literal,
+    Optional,
+    TYPE_CHECKING,
+    Union,
+)
 
 import requests
 
@@ -116,10 +126,10 @@ def delete_artifacts(self, **kwargs: Any) -> None:
     def artifacts(
         self,
         streamed: bool = False,
-        action: Optional[Callable[..., Any]] = None,
+        action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Get the job artifacts.
 
         Args:
@@ -152,10 +162,10 @@ def artifact(
         self,
         path: str,
         streamed: bool = False,
-        action: Optional[Callable[..., Any]] = None,
+        action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Get a single artifact file from within the job's artifacts archive.
 
         Args:
diff --git a/gitlab/v4/objects/packages.py b/gitlab/v4/objects/packages.py
@@ -5,7 +5,16 @@
 """
 
 from pathlib import Path
-from typing import Any, Callable, cast, Optional, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Callable,
+    cast,
+    Iterator,
+    Literal,
+    Optional,
+    TYPE_CHECKING,
+    Union,
+)
 
 import requests
 
@@ -103,10 +112,10 @@ def download(
         package_version: str,
         file_name: str,
         streamed: bool = False,
-        action: Optional[Callable] = None,
+        action: Optional[Union[Callable, Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Download a generic package.
 
         Args:
diff --git a/gitlab/v4/objects/projects.py b/gitlab/v4/objects/projects.py
@@ -1,4 +1,15 @@
-from typing import Any, Callable, cast, Dict, List, Optional, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Callable,
+    cast,
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    TYPE_CHECKING,
+    Union,
+)
 
 import requests
 
@@ -457,10 +468,10 @@ def snapshot(
         self,
         wiki: bool = False,
         streamed: bool = False,
-        action: Optional[Callable] = None,
+        action: Optional[Union[Callable, Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Return a snapshot of the repository.
 
         Args:
@@ -562,7 +573,7 @@ def artifact(
         self,
         *args: Any,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         utils.warn(
             message=(
                 "The project.artifact() method is deprecated and will be "
diff --git a/gitlab/v4/objects/repositories.py b/gitlab/v4/objects/repositories.py
@@ -3,7 +3,17 @@
 
 Currently this module only contains repository-related methods for projects.
 """
-from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    TYPE_CHECKING,
+    Union,
+)
 
 import requests
 
@@ -107,10 +117,10 @@ def repository_raw_blob(
         self,
         sha: str,
         streamed: bool = False,
-        action: Optional[Callable[..., Any]] = None,
+        action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Return the raw file contents for a blob.
 
         Args:
@@ -192,11 +202,11 @@ def repository_archive(
         self,
         sha: str = None,
         streamed: bool = False,
-        action: Optional[Callable[..., Any]] = None,
+        action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         format: Optional[str] = None,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Return an archive of the repository.
 
         Args:
diff --git a/gitlab/v4/objects/snippets.py b/gitlab/v4/objects/snippets.py
@@ -1,4 +1,14 @@
-from typing import Any, Callable, cast, List, Optional, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Callable,
+    cast,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    TYPE_CHECKING,
+    Union,
+)
 
 import requests
 
@@ -28,10 +38,10 @@ class Snippet(UserAgentDetailMixin, SaveMixin, ObjectDeleteMixin, RESTObject):
     def content(
         self,
         streamed: bool = False,
-        action: Optional[Callable[..., Any]] = None,
+        action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Return the content of a snippet.
 
         Args:
@@ -102,10 +112,10 @@ class ProjectSnippet(UserAgentDetailMixin, SaveMixin, ObjectDeleteMixin, RESTObj
     def content(
         self,
         streamed: bool = False,
-        action: Optional[Callable[..., Any]] = None,
+        action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
         chunk_size: int = 1024,
         **kwargs: Any,
-    ) -> Optional[bytes]:
+    ) -> Optional[Union[bytes, Iterator[Any]]]:
         """Return the content of a snippet.
 
         Args:
