feat(packages): Allow uploading bytes and files

ernestask · ernestask · commit a2beef324c7c · 2023-08-03T11:57:13.000+03:00
This commit adds a keyword argument to GenericPackageManager.upload() to allow uploading bytes and file-like objects to the generic package registry. That necessitates changing file path to be a keyword argument as well, which then cascades into a whole slew of checks to not allow passing both and to not allow uploading file-like objects as JSON data. Closes #1815
diff --git a/gitlab/_backends/requests_backend.py b/gitlab/_backends/requests_backend.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Any, Dict, Optional, TYPE_CHECKING, Union
+from typing import Any, BinaryIO, Dict, Optional, TYPE_CHECKING, Union
 
 import requests
 from requests.structures import CaseInsensitiveDict
@@ -63,7 +63,7 @@ def client(self) -> requests.Session:
     @staticmethod
     def prepare_send_data(
         files: Optional[Dict[str, Any]] = None,
-        post_data: Optional[Union[Dict[str, Any], bytes]] = None,
+        post_data: Optional[Union[Dict[str, Any], bytes, BinaryIO]] = None,
         raw: bool = False,
     ) -> SendData:
         if files:
@@ -90,6 +90,9 @@ def prepare_send_data(
         if raw and post_data:
             return SendData(data=post_data, content_type="application/octet-stream")
 
+        if TYPE_CHECKING:
+            assert not isinstance(post_data, BinaryIO)
+
         return SendData(json=post_data, content_type="application/json")
 
     def http_request(
diff --git a/gitlab/client.py b/gitlab/client.py
@@ -3,7 +3,18 @@
 import os
 import re
 import time
-from typing import Any, cast, Dict, List, Optional, Tuple, Type, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    BinaryIO,
+    cast,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TYPE_CHECKING,
+    Union,
+)
 from urllib import parse
 
 import requests
@@ -644,7 +655,7 @@ def http_request(
         verb: str,
         path: str,
         query_data: Optional[Dict[str, Any]] = None,
-        post_data: Optional[Union[Dict[str, Any], bytes]] = None,
+        post_data: Optional[Union[Dict[str, Any], bytes, BinaryIO]] = None,
         raw: bool = False,
         streamed: bool = False,
         files: Optional[Dict[str, Any]] = None,
@@ -1045,7 +1056,7 @@ def http_put(
         self,
         path: str,
         query_data: Optional[Dict[str, Any]] = None,
-        post_data: Optional[Union[Dict[str, Any], bytes]] = None,
+        post_data: Optional[Union[Dict[str, Any], bytes, BinaryIO]] = None,
         raw: bool = False,
         files: Optional[Dict[str, Any]] = None,
         **kwargs: Any,
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, Iterator, Optional, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    BinaryIO,
+    Callable,
+    cast,
+    Iterator,
+    Optional,
+    TYPE_CHECKING,
+    Union,
+)
 
 import requests
 
@@ -46,8 +55,9 @@ def upload(
         package_name: str,
         package_version: str,
         file_name: str,
-        path: Union[str, Path],
+        path: Optional[Union[str, Path]] = None,
         select: Optional[str] = None,
+        data: Optional[Union[bytes, BinaryIO]] = None,
         **kwargs: Any,
     ) -> GenericPackage:
         """Upload a file as a generic package.
@@ -72,11 +82,25 @@ def upload(
         https://docs.gitlab.com/ee/user/packages/generic_packages/
         """
 
-        try:
-            with open(path, "rb") as f:
-                file_data = f.read()
-        except OSError as e:
-            raise exc.GitlabUploadError(f"Failed to read package file {path}") from e
+        if path is None and data is None:
+            raise exc.GitlabUploadError("No file contents or path specified")
+
+        if path is not None and data is not None:
+            raise exc.GitlabUploadError("File contents and file path specified")
+
+        file_data: Optional[Union[bytes, BinaryIO]] = data
+
+        if not file_data:
+            if TYPE_CHECKING:
+                assert path is not None
+
+            try:
+                with open(path, "rb") as f:
+                    file_data = f.read()
+            except OSError as e:
+                raise exc.GitlabUploadError(
+                    f"Failed to read package file {path}"
+                ) from e
 
         url = f"{self._computed_path}/{package_name}/{package_version}/{file_name}"
         query_data = {} if select is None else {"select": select}
diff --git a/tests/functional/api/test_packages.py b/tests/functional/api/test_packages.py
@@ -38,6 +38,38 @@ def test_upload_generic_package(tmp_path, project):
     assert package.message == "201 Created"
 
 
+def test_download_generic_package_bytes(tmp_path, project):
+    path = tmp_path / file_name
+
+    path.write_text(file_content)
+
+    package = project.generic_packages.upload(
+        package_name=package_name,
+        package_version=package_version,
+        file_name=file_name,
+        data=path.read_bytes(),
+    )
+
+    assert isinstance(package, GenericPackage)
+    assert package.message == "201 Created"
+
+
+def test_download_generic_package_file(tmp_path, project):
+    path = tmp_path / file_name
+
+    path.write_text(file_content)
+
+    package = project.generic_packages.upload(
+        package_name=package_name,
+        package_version=package_version,
+        file_name=file_name,
+        data=path.open(mode="rb"),
+    )
+
+    assert isinstance(package, GenericPackage)
+    assert package.message == "201 Created"
+
+
 def test_upload_generic_package_select(tmp_path, project):
     path = tmp_path / file_name2
     path.write_text(file_content)
diff --git a/tests/unit/objects/test_packages.py b/tests/unit/objects/test_packages.py
@@ -287,6 +287,36 @@ def test_upload_generic_package(tmp_path, project, resp_upload_generic_package):
     assert isinstance(package, GenericPackage)
 
 
+def test_upload_generic_package_bytes(tmp_path, project, resp_upload_generic_package):
+    path = tmp_path / file_name
+
+    path.write_text(file_content, encoding="utf-8")
+
+    package = project.generic_packages.upload(
+        package_name=package_name,
+        package_version=package_version,
+        file_name=file_name,
+        data=path.read_bytes(),
+    )
+
+    assert isinstance(package, GenericPackage)
+
+
+def test_upload_generic_package_file(tmp_path, project, resp_upload_generic_package):
+    path = tmp_path / file_name
+
+    path.write_text(file_content, encoding="utf-8")
+
+    package = project.generic_packages.upload(
+        package_name=package_name,
+        package_version=package_version,
+        file_name=file_name,
+        data=path.open(mode="rb"),
+    )
+
+    assert isinstance(package, GenericPackage)
+
+
 def test_download_generic_package(project, resp_download_generic_package):
     package = project.generic_packages.download(
         package_name=package_name,
