Skip to content

Commit

Permalink
Add method to download an attachment
Browse files Browse the repository at this point in the history
The browser_download_url property of attachments can be confusing, as it
is a full URL and `allspice.requests_get_raw` doesn't work with absolute
URLs. The new `Attachment.download_to_file` method provides a "blessed"
way to download the attachment.
  • Loading branch information
shrik450 committed Sep 16, 2024
1 parent a9efd29 commit 28c72a6
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 0 deletions.
22 changes: 22 additions & 0 deletions allspice/apiobject.py
Original file line number Diff line number Diff line change
Expand Up @@ -1531,6 +1531,28 @@ def __eq__(self, other):
def __hash__(self):
return hash(self.uuid)

def download_to_file(self, io: IO):
"""
Download the raw, binary data of this Attachment to a file-like object.
Example:
with open("my_file.zip", "wb") as f:
attachment.download_to_file(f)
:param io: The file-like object to write the data to.
"""

response = self.allspice_client.requests.get(
self.browser_download_url,
headers=self.allspice_client.headers,
stream=True,
)
# 4kb chunks
for chunk in response.iter_content(chunk_size=4096):
if chunk:
io.write(chunk)


class Comment(ApiObject):
assets: List[Union[Any, Dict[str, Union[int, str]]]]
Expand Down
21 changes: 21 additions & 0 deletions tests/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -627,6 +627,27 @@ def test_get_issue_attachments(instance):
assert attachments[0].name == "requirements.txt"


def test_download_issue_attachment(instance, tmp_path):
org = Organization.request(instance, test_org)
repo = Repository.request(instance, org.username, test_repo)
issue = repo.get_issues()[0]
comment = issue.get_comments()[0]
attachment = comment.get_attachments()[0]

filename = uuid.uuid4().hex[:8] + ".txt"
filepath = tmp_path / filename
with open(filepath, "wb") as f:
attachment.download_to_file(f)

with open(filepath, "r") as actual_f:
with open("requirements.txt", "r") as expected_f:
attachment_content = actual_f.read()
assert len(attachment_content) > 0

expected_content = expected_f.read()
assert expected_content == attachment_content


def test_edit_issue_attachment(instance):
org = Organization.request(instance, test_org)
repo = Repository.request(instance, org.username, test_repo)
Expand Down

0 comments on commit 28c72a6

Please sign in to comment.