diff --git a/docs/gl_objects/pipelines_and_jobs.rst b/docs/gl_objects/pipelines_and_jobs.rst index a05d968a4..f0bdd3a68 100644 --- a/docs/gl_objects/pipelines_and_jobs.rst +++ b/docs/gl_objects/pipelines_and_jobs.rst @@ -274,6 +274,19 @@ You can also directly stream the output into a file, and unzip it afterwards:: subprocess.run(["unzip", "-bo", zipfn]) os.unlink(zipfn) +Or, you can also use the underlying response iterator directly:: + + artifact_bytes_iterator = build_or_job.artifacts(iterator=True) + +This can be used with frameworks that expect an iterator (such as FastAPI/Starlette's +``StreamingResponse``) to forward a download from GitLab without having to download +the entire content server-side first:: + + @app.get("/download_artifact") + def download_artifact(): + artifact_bytes_iterator = build_or_job.artifacts(iterator=True) + return StreamingResponse(artifact_bytes_iterator, media_type="application/zip") + Delete all artifacts of a project that can be deleted:: project.artifacts.delete() diff --git a/gitlab/mixins.py b/gitlab/mixins.py index 71ba8210c..14542e0a6 100644 --- a/gitlab/mixins.py +++ b/gitlab/mixins.py @@ -20,6 +20,7 @@ Any, Callable, Dict, + Iterator, List, Optional, Tuple, @@ -614,16 +615,19 @@ class DownloadMixin(_RestObjectBase): def download( self, streamed: bool = False, + iterator: bool = False, action: Optional[Callable] = None, chunk_size: int = 1024, **kwargs: Any, - ) -> Optional[bytes]: + ) -> Optional[Union[bytes, Iterator[Any]]]: """Download the archive of a resource export. Args: streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -642,7 +646,7 @@ def download( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) class SubscribableMixin(_RestObjectBase): diff --git a/gitlab/utils.py b/gitlab/utils.py index bab670584..6acb86160 100644 --- 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, Tuple, Type, Union +from typing import Any, Callable, Dict, Iterator, Optional, Tuple, Type, Union import requests @@ -34,9 +34,13 @@ def __call__(self, chunk: Any) -> None: def response_content( response: requests.Response, streamed: bool, + iterator: bool, action: Optional[Callable], chunk_size: int, -) -> Optional[bytes]: +) -> Optional[Union[bytes, Iterator[Any]]]: + if iterator: + return response.iter_content(chunk_size=chunk_size) + if streamed is False: return response.content diff --git a/gitlab/v4/cli.py b/gitlab/v4/cli.py index 2b0d4ce72..ba2e788b7 100644 --- a/gitlab/v4/cli.py +++ b/gitlab/v4/cli.py @@ -127,6 +127,7 @@ def do_project_export_download(self) -> None: data = export_status.download() if TYPE_CHECKING: assert data is not None + assert isinstance(data, bytes) sys.stdout.buffer.write(data) except Exception as e: # pragma: no cover, cli.die is unit-tested diff --git a/gitlab/v4/objects/artifacts.py b/gitlab/v4/objects/artifacts.py index 541e5e2f4..f5f106d8b 100644 --- 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, Optional, TYPE_CHECKING, Union import requests @@ -40,10 +40,14 @@ def __call__( ), category=DeprecationWarning, ) - return self.download( + data = self.download( *args, **kwargs, ) + if TYPE_CHECKING: + assert data is not None + assert isinstance(data, bytes) + return data @exc.on_http_error(exc.GitlabDeleteError) def delete(self, **kwargs: Any) -> None: @@ -71,10 +75,11 @@ def download( ref_name: str, job: str, streamed: bool = False, + iterator: bool = False, action: Optional[Callable] = 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: @@ -85,6 +90,8 @@ def download( streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -103,7 +110,7 @@ def download( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) @cli.register_custom_action( "ProjectArtifactManager", ("ref_name", "artifact_path", "job") @@ -115,10 +122,11 @@ def raw( artifact_path: str, job: str, streamed: bool = False, + iterator: bool = False, action: Optional[Callable] = 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. @@ -130,6 +138,8 @@ def raw( streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -148,4 +158,4 @@ def raw( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) diff --git a/gitlab/v4/objects/files.py b/gitlab/v4/objects/files.py index aa86704c9..2fd79fd54 100644 --- a/gitlab/v4/objects/files.py +++ b/gitlab/v4/objects/files.py @@ -1,5 +1,15 @@ import base64 -from typing import Any, Callable, cast, Dict, List, Optional, TYPE_CHECKING +from typing import ( + Any, + Callable, + cast, + Dict, + Iterator, + List, + Optional, + TYPE_CHECKING, + Union, +) import requests @@ -220,10 +230,11 @@ def raw( file_path: str, ref: str, streamed: bool = False, + iterator: bool = False, action: Optional[Callable[..., Any]] = None, chunk_size: int = 1024, **kwargs: Any, - ) -> Optional[bytes]: + ) -> Optional[Union[bytes, Iterator[Any]]]: """Return the content of a file for a commit. Args: @@ -232,6 +243,8 @@ def raw( streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -252,7 +265,7 @@ def raw( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) @cli.register_custom_action("ProjectFileManager", ("file_path", "ref")) @exc.on_http_error(exc.GitlabListError) diff --git a/gitlab/v4/objects/jobs.py b/gitlab/v4/objects/jobs.py index fbcb1fd40..850227725 100644 --- a/gitlab/v4/objects/jobs.py +++ b/gitlab/v4/objects/jobs.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, cast, Dict, Optional, TYPE_CHECKING, Union +from typing import Any, Callable, cast, Dict, Iterator, Optional, TYPE_CHECKING, Union import requests @@ -116,16 +116,19 @@ def delete_artifacts(self, **kwargs: Any) -> None: def artifacts( self, streamed: bool = False, + iterator: bool = False, action: Optional[Callable[..., Any]] = None, chunk_size: int = 1024, **kwargs: Any, - ) -> Optional[bytes]: + ) -> Optional[Union[bytes, Iterator[Any]]]: """Get the job artifacts. Args: streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -144,7 +147,7 @@ def artifacts( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) @cli.register_custom_action("ProjectJob") @exc.on_http_error(exc.GitlabGetError) @@ -152,10 +155,11 @@ def artifact( self, path: str, streamed: bool = False, + iterator: bool = False, action: Optional[Callable[..., Any]] = 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: @@ -163,6 +167,8 @@ def artifact( streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -181,13 +187,14 @@ def artifact( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) @cli.register_custom_action("ProjectJob") @exc.on_http_error(exc.GitlabGetError) def trace( self, streamed: bool = False, + iterator: bool = False, action: Optional[Callable[..., Any]] = None, chunk_size: int = 1024, **kwargs: Any, @@ -198,6 +205,8 @@ def trace( streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -216,7 +225,9 @@ def trace( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return_value = utils.response_content(result, streamed, action, chunk_size) + return_value = utils.response_content( + result, streamed, iterator, action, chunk_size + ) if TYPE_CHECKING: assert isinstance(return_value, dict) return return_value diff --git a/gitlab/v4/objects/packages.py b/gitlab/v4/objects/packages.py index 882cb1a5a..a82080167 100644 --- a/gitlab/v4/objects/packages.py +++ b/gitlab/v4/objects/packages.py @@ -5,7 +5,7 @@ """ from pathlib import Path -from typing import Any, Callable, cast, Optional, TYPE_CHECKING, Union +from typing import Any, Callable, cast, Iterator, Optional, TYPE_CHECKING, Union import requests @@ -103,10 +103,11 @@ def download( package_version: str, file_name: str, streamed: bool = False, + iterator: bool = False, action: Optional[Callable] = None, chunk_size: int = 1024, **kwargs: Any, - ) -> Optional[bytes]: + ) -> Optional[Union[bytes, Iterator[Any]]]: """Download a generic package. Args: @@ -116,6 +117,8 @@ def download( streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -132,7 +135,7 @@ def download( result = self.gitlab.http_get(path, streamed=streamed, raw=True, **kwargs) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) class GroupPackage(RESTObject): diff --git a/gitlab/v4/objects/projects.py b/gitlab/v4/objects/projects.py index f32cf2257..fc2aac3d5 100644 --- a/gitlab/v4/objects/projects.py +++ b/gitlab/v4/objects/projects.py @@ -1,4 +1,14 @@ -from typing import Any, Callable, cast, Dict, List, Optional, TYPE_CHECKING, Union +from typing import ( + Any, + Callable, + cast, + Dict, + Iterator, + List, + Optional, + TYPE_CHECKING, + Union, +) import requests @@ -466,10 +476,11 @@ def snapshot( self, wiki: bool = False, streamed: bool = False, + iterator: bool = False, action: Optional[Callable] = None, chunk_size: int = 1024, **kwargs: Any, - ) -> Optional[bytes]: + ) -> Optional[Union[bytes, Iterator[Any]]]: """Return a snapshot of the repository. Args: @@ -477,6 +488,8 @@ def snapshot( streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment. + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -495,7 +508,7 @@ def snapshot( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) @cli.register_custom_action("Project", ("scope", "search")) @exc.on_http_error(exc.GitlabSearchError) @@ -579,7 +592,11 @@ def artifact( ), category=DeprecationWarning, ) - return self.artifacts.raw(*args, **kwargs) + data = self.artifacts.raw(*args, **kwargs) + if TYPE_CHECKING: + assert data is not None + assert isinstance(data, bytes) + return data class ProjectManager(CRUDMixin, RESTManager): diff --git a/gitlab/v4/objects/repositories.py b/gitlab/v4/objects/repositories.py index 5826d9d83..1f10473aa 100644 --- a/gitlab/v4/objects/repositories.py +++ b/gitlab/v4/objects/repositories.py @@ -3,7 +3,7 @@ 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, Optional, TYPE_CHECKING, Union import requests @@ -107,10 +107,11 @@ def repository_raw_blob( self, sha: str, streamed: bool = False, + iterator: bool = False, action: Optional[Callable[..., Any]] = None, chunk_size: int = 1024, **kwargs: Any, - ) -> Optional[bytes]: + ) -> Optional[Union[bytes, Iterator[Any]]]: """Return the raw file contents for a blob. Args: @@ -118,6 +119,8 @@ def repository_raw_blob( streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -136,7 +139,7 @@ def repository_raw_blob( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) @cli.register_custom_action("Project", ("from_", "to")) @exc.on_http_error(exc.GitlabGetError) @@ -192,11 +195,12 @@ def repository_archive( self, sha: str = None, streamed: bool = False, + iterator: bool = False, action: Optional[Callable[..., Any]] = 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: @@ -204,6 +208,8 @@ def repository_archive( streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -228,7 +234,7 @@ def repository_archive( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) @cli.register_custom_action("Project") @exc.on_http_error(exc.GitlabDeleteError) diff --git a/gitlab/v4/objects/snippets.py b/gitlab/v4/objects/snippets.py index 597a3aaf0..aa46c7747 100644 --- a/gitlab/v4/objects/snippets.py +++ b/gitlab/v4/objects/snippets.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, cast, List, Optional, TYPE_CHECKING, Union +from typing import Any, Callable, cast, Iterator, List, Optional, TYPE_CHECKING, Union import requests @@ -29,16 +29,19 @@ class Snippet(UserAgentDetailMixin, SaveMixin, ObjectDeleteMixin, RESTObject): def content( self, streamed: bool = False, + iterator: bool = False, action: Optional[Callable[..., Any]] = None, chunk_size: int = 1024, **kwargs: Any, - ) -> Optional[bytes]: + ) -> Optional[Union[bytes, Iterator[Any]]]: """Return the content of a snippet. Args: streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment. + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -57,7 +60,7 @@ def content( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) class SnippetManager(CRUDMixin, RESTManager): @@ -103,16 +106,19 @@ class ProjectSnippet(UserAgentDetailMixin, SaveMixin, ObjectDeleteMixin, RESTObj def content( self, streamed: bool = False, + iterator: bool = False, action: Optional[Callable[..., Any]] = None, chunk_size: int = 1024, **kwargs: Any, - ) -> Optional[bytes]: + ) -> Optional[Union[bytes, Iterator[Any]]]: """Return the content of a snippet. Args: streamed: If True the data will be processed by chunks of `chunk_size` and each chunk is passed to `action` for treatment. + iterator: If True directly return the underlying response + iterator action: Callable responsible of dealing with chunk of data chunk_size: Size of each chunk @@ -131,7 +137,7 @@ def content( ) if TYPE_CHECKING: assert isinstance(result, requests.Response) - return utils.response_content(result, streamed, action, chunk_size) + return utils.response_content(result, streamed, iterator, action, chunk_size) class ProjectSnippetManager(CRUDMixin, RESTManager): diff --git a/tests/functional/api/test_packages.py b/tests/functional/api/test_packages.py index 64b57b827..9f06439b4 100644 --- a/tests/functional/api/test_packages.py +++ b/tests/functional/api/test_packages.py @@ -3,6 +3,8 @@ https://docs.gitlab.com/ce/api/packages.html https://docs.gitlab.com/ee/user/packages/generic_packages """ +from collections.abc import Iterator + from gitlab.v4.objects import GenericPackage package_name = "hello-world" @@ -46,6 +48,24 @@ def test_download_generic_package(project): assert package.decode("utf-8") == file_content +def test_stream_generic_package(project): + bytes_iterator = project.generic_packages.download( + package_name=package_name, + package_version=package_version, + file_name=file_name, + iterator=True, + ) + + assert isinstance(bytes_iterator, Iterator) + + package = bytes() + for chunk in bytes_iterator: + package += chunk + + assert isinstance(package, bytes) + assert package.decode("utf-8") == file_content + + def test_download_generic_package_to_file(tmp_path, project): path = tmp_path / file_name @@ -60,3 +80,21 @@ def test_download_generic_package_to_file(tmp_path, project): with open(path, "r") as f: assert f.read() == file_content + + +def test_stream_generic_package_to_file(tmp_path, project): + path = tmp_path / file_name + + bytes_iterator = project.generic_packages.download( + package_name=package_name, + package_version=package_version, + file_name=file_name, + iterator=True, + ) + + with open(path, "wb") as f: + for chunk in bytes_iterator: + f.write(chunk) + + with open(path, "r") as f: + assert f.read() == file_content diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index c5f9f931d..74b48ae31 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -36,7 +36,9 @@ def test_response_content(capsys): ) resp = requests.get("https://example.com", stream=True) - utils.response_content(resp, streamed=True, action=None, chunk_size=1024) + utils.response_content( + resp, streamed=True, iterator=False, action=None, chunk_size=1024 + ) captured = capsys.readouterr() assert "test" in captured.out