Skip to content

Commit 4f9807f

Browse files
TCatshoekTCatshoek
committed
feat(downloads): allow streaming downloads access to response iterator
Allow access to the underlying response iterator when downloading in streaming mode by specifying action="iterator" Update type annotations to support this change
1 parent 5370979 commit 4f9807f

File tree

10 files changed

+106
-38
lines changed

10 files changed

+106
-38
lines changed

‎gitlab/mixins.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,9 @@
2020
Any,
2121
Callable,
2222
Dict,
23+
Iterator,
2324
List,
25+
Literal,
2426
Optional,
2527
Tuple,
2628
Type,
@@ -657,10 +659,10 @@ class DownloadMixin(_RestObjectBase):
657659
def download(
658660
self,
659661
streamed: bool = False,
660-
action: Optional[Callable] = None,
662+
action: Optional[Union[Callable, Literal["iterator"]]] = None,
661663
chunk_size: int = 1024,
662664
**kwargs: Any,
663-
) -> Optional[bytes]:
665+
) -> Optional[Union[bytes, Iterator[Any]]]:
664666
"""Download the archive of a resource export.
665667
666668
Args:

‎gitlab/utils.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import traceback
2020
import urllib.parse
2121
import warnings
22-
from typing import Any, Callable, Dict, Optional, Type, Union
22+
from typing import Any, Callable, Dict, Iterator, Literal, Optional, Type, Union
2323

2424
import requests
2525

@@ -32,15 +32,18 @@ def __call__(self, chunk: Any) -> None:
3232
def response_content(
3333
response: requests.Response,
3434
streamed: bool,
35-
action: Optional[Callable],
35+
action: Optional[Union[Callable, Literal["iterator"]]],
3636
chunk_size: int,
37-
) -> Optional[bytes]:
37+
) -> Optional[Union[bytes, Iterator[Any]]]:
3838
if streamed is False:
3939
return response.content
4040

4141
if action is None:
4242
action = _StdoutStream()
4343

44+
if action == "iterator":
45+
return response.iter_content(chunk_size=chunk_size)
46+
4447
for chunk in response.iter_content(chunk_size=chunk_size):
4548
if chunk:
4649
action(chunk)

‎gitlab/v4/cli.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import argparse
2020
import operator
2121
import sys
22-
from typing import Any, Dict, List, Optional, Type, TYPE_CHECKING, Union
22+
from typing import Any, Dict, Iterator, List, Optional, Type, TYPE_CHECKING, Union
2323

2424
import gitlab
2525
import gitlab.base
@@ -123,6 +123,7 @@ def do_project_export_download(self) -> None:
123123
data = export_status.download()
124124
if TYPE_CHECKING:
125125
assert data is not None
126+
assert not isinstance(data, Iterator)
126127
sys.stdout.buffer.write(data)
127128

128129
except Exception as e:

‎gitlab/v4/objects/artifacts.py

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
GitLab API:
33
https://docs.gitlab.com/ee/api/job_artifacts.html
44
"""
5-
from typing import Any, Callable, Optional, TYPE_CHECKING
5+
from typing import Any, Callable, Iterator, Literal, Optional, TYPE_CHECKING, Union
66

77
import requests
88

@@ -32,7 +32,7 @@ def __call__(
3232
self,
3333
*args: Any,
3434
**kwargs: Any,
35-
) -> Optional[bytes]:
35+
) -> Optional[Union[bytes, Iterator[Any]]]:
3636
utils.warn(
3737
message=(
3838
"The project.artifacts() method is deprecated and will be removed in a "
@@ -71,10 +71,10 @@ def download(
7171
ref_name: str,
7272
job: str,
7373
streamed: bool = False,
74-
action: Optional[Callable] = None,
74+
action: Optional[Union[Callable, Literal["iterator"]]] = None,
7575
chunk_size: int = 1024,
7676
**kwargs: Any,
77-
) -> Optional[bytes]:
77+
) -> Optional[Union[bytes, Iterator[Any]]]:
7878
"""Get the job artifacts archive from a specific tag or branch.
7979
8080
Args:
@@ -86,7 +86,8 @@ def download(
8686
`chunk_size` and each chunk is passed to `action` for
8787
treatment
8888
action: Callable responsible of dealing with chunk of
89-
data
89+
data. May also be the string "iterator" to directly return
90+
the response iterator
9091
chunk_size: Size of each chunk
9192
**kwargs: Extra options to send to the server (e.g. sudo)
9293
@@ -115,10 +116,10 @@ def raw(
115116
artifact_path: str,
116117
job: str,
117118
streamed: bool = False,
118-
action: Optional[Callable] = None,
119+
action: Optional[Union[Callable, Literal["iterator"]]] = None,
119120
chunk_size: int = 1024,
120121
**kwargs: Any,
121-
) -> Optional[bytes]:
122+
) -> Optional[Union[bytes, Iterator[Any]]]:
122123
"""Download a single artifact file from a specific tag or branch from
123124
within the job's artifacts archive.
124125

‎gitlab/v4/objects/files.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,16 @@
11
import base64
2-
from typing import Any, Callable, cast, Dict, List, Optional, TYPE_CHECKING
2+
from typing import (
3+
Any,
4+
Callable,
5+
cast,
6+
Dict,
7+
Iterator,
8+
List,
9+
Literal,
10+
Optional,
11+
TYPE_CHECKING,
12+
Union,
13+
)
314

415
import requests
516

@@ -217,10 +228,10 @@ def raw(
217228
file_path: str,
218229
ref: str,
219230
streamed: bool = False,
220-
action: Optional[Callable[..., Any]] = None,
231+
action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
221232
chunk_size: int = 1024,
222233
**kwargs: Any,
223-
) -> Optional[bytes]:
234+
) -> Optional[Union[bytes, Iterator[Any]]]:
224235
"""Return the content of a file for a commit.
225236
226237
Args:

‎gitlab/v4/objects/jobs.py

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,14 @@
1-
from typing import Any, Callable, cast, Dict, Optional, TYPE_CHECKING, Union
1+
from typing import (
2+
Any,
3+
Callable,
4+
cast,
5+
Dict,
6+
Iterator,
7+
Literal,
8+
Optional,
9+
TYPE_CHECKING,
10+
Union,
11+
)
212

313
import requests
414

@@ -116,10 +126,10 @@ def delete_artifacts(self, **kwargs: Any) -> None:
116126
def artifacts(
117127
self,
118128
streamed: bool = False,
119-
action: Optional[Callable[..., Any]] = None,
129+
action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
120130
chunk_size: int = 1024,
121131
**kwargs: Any,
122-
) -> Optional[bytes]:
132+
) -> Optional[Union[bytes, Iterator[Any]]]:
123133
"""Get the job artifacts.
124134
125135
Args:
@@ -152,10 +162,10 @@ def artifact(
152162
self,
153163
path: str,
154164
streamed: bool = False,
155-
action: Optional[Callable[..., Any]] = None,
165+
action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
156166
chunk_size: int = 1024,
157167
**kwargs: Any,
158-
) -> Optional[bytes]:
168+
) -> Optional[Union[bytes, Iterator[Any]]]:
159169
"""Get a single artifact file from within the job's artifacts archive.
160170
161171
Args:

‎gitlab/v4/objects/packages.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,16 @@
55
"""
66

77
from pathlib import Path
8-
from typing import Any, Callable, cast, Optional, TYPE_CHECKING, Union
8+
from typing import (
9+
Any,
10+
Callable,
11+
cast,
12+
Iterator,
13+
Literal,
14+
Optional,
15+
TYPE_CHECKING,
16+
Union,
17+
)
918

1019
import requests
1120

@@ -103,10 +112,10 @@ def download(
103112
package_version: str,
104113
file_name: str,
105114
streamed: bool = False,
106-
action: Optional[Callable] = None,
115+
action: Optional[Union[Callable, Literal["iterator"]]] = None,
107116
chunk_size: int = 1024,
108117
**kwargs: Any,
109-
) -> Optional[bytes]:
118+
) -> Optional[Union[bytes, Iterator[Any]]]:
110119
"""Download a generic package.
111120
112121
Args:

‎gitlab/v4/objects/projects.py

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,15 @@
1-
from typing import Any, Callable, cast, Dict, List, Optional, TYPE_CHECKING, Union
1+
from typing import (
2+
Any,
3+
Callable,
4+
cast,
5+
Dict,
6+
Iterator,
7+
List,
8+
Literal,
9+
Optional,
10+
TYPE_CHECKING,
11+
Union,
12+
)
213

314
import requests
415

@@ -457,10 +468,10 @@ def snapshot(
457468
self,
458469
wiki: bool = False,
459470
streamed: bool = False,
460-
action: Optional[Callable] = None,
471+
action: Optional[Union[Callable, Literal["iterator"]]] = None,
461472
chunk_size: int = 1024,
462473
**kwargs: Any,
463-
) -> Optional[bytes]:
474+
) -> Optional[Union[bytes, Iterator[Any]]]:
464475
"""Return a snapshot of the repository.
465476
466477
Args:
@@ -562,7 +573,7 @@ def artifact(
562573
self,
563574
*args: Any,
564575
**kwargs: Any,
565-
) -> Optional[bytes]:
576+
) -> Optional[Union[bytes, Iterator[Any]]]:
566577
utils.warn(
567578
message=(
568579
"The project.artifact() method is deprecated and will be "

‎gitlab/v4/objects/repositories.py

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,17 @@
33
44
Currently this module only contains repository-related methods for projects.
55
"""
6-
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING, Union
6+
from typing import (
7+
Any,
8+
Callable,
9+
Dict,
10+
Iterator,
11+
List,
12+
Literal,
13+
Optional,
14+
TYPE_CHECKING,
15+
Union,
16+
)
717

818
import requests
919

@@ -107,10 +117,10 @@ def repository_raw_blob(
107117
self,
108118
sha: str,
109119
streamed: bool = False,
110-
action: Optional[Callable[..., Any]] = None,
120+
action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
111121
chunk_size: int = 1024,
112122
**kwargs: Any,
113-
) -> Optional[bytes]:
123+
) -> Optional[Union[bytes, Iterator[Any]]]:
114124
"""Return the raw file contents for a blob.
115125
116126
Args:
@@ -192,11 +202,11 @@ def repository_archive(
192202
self,
193203
sha: str = None,
194204
streamed: bool = False,
195-
action: Optional[Callable[..., Any]] = None,
205+
action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
196206
chunk_size: int = 1024,
197207
format: Optional[str] = None,
198208
**kwargs: Any,
199-
) -> Optional[bytes]:
209+
) -> Optional[Union[bytes, Iterator[Any]]]:
200210
"""Return an archive of the repository.
201211
202212
Args:

‎gitlab/v4/objects/snippets.py

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,14 @@
1-
from typing import Any, Callable, cast, List, Optional, TYPE_CHECKING, Union
1+
from typing import (
2+
Any,
3+
Callable,
4+
cast,
5+
Iterator,
6+
List,
7+
Literal,
8+
Optional,
9+
TYPE_CHECKING,
10+
Union,
11+
)
212

313
import requests
414

@@ -28,10 +38,10 @@ class Snippet(UserAgentDetailMixin, SaveMixin, ObjectDeleteMixin, RESTObject):
2838
def content(
2939
self,
3040
streamed: bool = False,
31-
action: Optional[Callable[..., Any]] = None,
41+
action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
3242
chunk_size: int = 1024,
3343
**kwargs: Any,
34-
) -> Optional[bytes]:
44+
) -> Optional[Union[bytes, Iterator[Any]]]:
3545
"""Return the content of a snippet.
3646
3747
Args:
@@ -102,10 +112,10 @@ class ProjectSnippet(UserAgentDetailMixin, SaveMixin, ObjectDeleteMixin, RESTObj
102112
def content(
103113
self,
104114
streamed: bool = False,
105-
action: Optional[Callable[..., Any]] = None,
115+
action: Optional[Union[Callable[..., Any], Literal["iterator"]]] = None,
106116
chunk_size: int = 1024,
107117
**kwargs: Any,
108-
) -> Optional[bytes]:
118+
) -> Optional[Union[bytes, Iterator[Any]]]:
109119
"""Return the content of a snippet.
110120
111121
Args:

0 commit comments

Comments
 (0)