Repository: omarryhan/aiogoogle

Issue:
feature request: support for in memory upload/download media
Thanks a lot for this great library!

I would like to use it in a tool that performs automated backups of autodesk bim360 files (but it could be any other storage service) to a google drive folder.

As of now, the only way to use the drive method `files.create()` is to specify the `upload_file` parameter as the file path, but I would like to use a file object (the one downloaded from the other service) without writing it to memory first.

I see that the `fire_request` coroutine in `AiohttpSession.send` does the reading of the file via `aiofiles.open`, and then passes the file object (or `MultipartWriter` object) to `self.request`.

Do you think it could be a good addition to support my use case?

If so, I could try to make a PR for that, but I would like to know your thoughts on this (and tips on how to accomplish this in a clean way).

Discussion:

Patch:
diff --git a/aiogoogle/models.py b/aiogoogle/models.py
index 083fad4..f620a35 100644
--- a/aiogoogle/models.py
+++ b/aiogoogle/models.py
@@ -1,8 +1,7 @@
 from urllib.parse import urlparse, urlunparse, urlencode, parse_qs
 import pprint
 
-from .excs import HTTPError, AuthError
-
+from .excs import HTTPError, AuthError, ValidationError
 
 DEFAULT_DOWNLOAD_CHUNK_SIZE = 1024 * 1024
 DEFAULT_UPLOAD_CHUNK_SIZE = 1024 * 1024
@@ -20,17 +19,14 @@ class ResumableUpload:
         
         multipart (bool): True if this endpoint supports upload multipart media.
 
-        chunksize (int): Size of a chunk of bytes that a session should read at a time when uploading in multipart.
+        chunk_size (int): Size of a chunk of bytes that a session should read at a time when uploading in multipart.
 
     """
 
-    def __init__(self, file_path, multipart=None, chunk_size=None, upload_path=None):
-        self.file_path = file_path
+    def __init__(self, multipart=None, chunk_size=None, upload_path=None):
         self.upload_path = upload_path
         self.multipart = multipart
-        if chunk_size is None:
-            chunk_size = DEFAULT_UPLOAD_CHUNK_SIZE
-        self.chunk_size = chunk_size
+        self.chunk_size = chunk_size or DEFAULT_UPLOAD_CHUNK_SIZE
 
 
 class MediaUpload:
@@ -40,8 +36,8 @@ class MediaUpload:
 
     Arguments:
 
-        file_path (str): Full path of the file to be uploaded
-        
+        file_path_or_bytes (str, bytes): Full path or content of the file to be uploaded
+
         upload_path (str): The URI path to be used for upload. Should be used in conjunction with the rootURL property at the API-level.
         
         mime_range (list): list of MIME Media Ranges for acceptable media uploads to this method.
@@ -60,7 +56,7 @@ class MediaUpload:
 
     def __init__(
         self,
-        file_path,
+        file_path_or_bytes,
         upload_path=None,
         mime_range=None,
         max_size=None,
@@ -69,17 +65,48 @@ class MediaUpload:
         resumable=None,
         validate=True,
     ):
-        self.file_path = file_path
+        if isinstance(file_path_or_bytes, bytes):
+            self.file_body = file_path_or_bytes
+            self.file_path = None
+        else:
+            self.file_body = None
+            self.file_path = file_path_or_bytes
         self.upload_path = upload_path
         self.mime_range = mime_range
         self.max_size = max_size
         self.multipart = multipart
-        if chunk_size is None:
-            chunk_size = DEFAULT_UPLOAD_CHUNK_SIZE
-        self.chunk_size = chunk_size
+        self.chunk_size = chunk_size or DEFAULT_UPLOAD_CHUNK_SIZE
         self.resumable = resumable
         self.validate = validate
 
+    async def run_validation(self, size_func):
+        if self.validate and self.max_size:
+            size = await size_func(self.file_path) if self.file_path else len(self.file_body)
+            if size > self.max_size:
+                raise ValidationError(
+                    f'"{self}" has a size of {size / 1000}KB. '
+                    f'Max upload size for this endpoint is: '
+                    f'{self.max_size / 1000}KB.'
+                )
+
+    async def aiter_file(self, aiter_func):
+        if self.file_path:
+            async for chunk in aiter_func(self.file_path, self.chunk_size):
+                yield chunk
+        else:
+            async for chunk in self._aiter_body():
+                yield chunk
+
+    async def _aiter_body(self):
+        for x in range(0, len(self.file_body), self.chunk_size):
+            yield self.file_body[x:x + self.chunk_size]
+
+    async def read_file(self, read_func):
+        return self.file_body or await read_func(self.file_path)
+
+    def __str__(self):
+        return self.file_path or "File object"
+
 
 class MediaDownload:
     """
@@ -95,9 +122,7 @@ class MediaDownload:
 
     def __init__(self, file_path, chunk_size=None):
         self.file_path = file_path
-        if chunk_size is None:
-            chunk_size = DEFAULT_DOWNLOAD_CHUNK_SIZE
-        self.chunk_size = chunk_size
+        self.chunk_size = chunk_size or DEFAULT_DOWNLOAD_CHUNK_SIZE
 
 
 class Request:
@@ -157,10 +182,7 @@ class Request:
         self.method = method
         self.url = url
         self.batch_url = batch_url
-        if headers is None:
-            self.headers = {}
-        else:
-            self.headers = headers
+        self.headers = {} if headers is None else headers
         self.data = data
         self.json = json
         self.media_upload = media_upload
diff --git a/aiogoogle/resource.py b/aiogoogle/resource.py
index c7916a8..d33e441 100644
--- a/aiogoogle/resource.py
+++ b/aiogoogle/resource.py
@@ -612,18 +612,17 @@ class Method:
             return base_url + self["path"]
 
     def _build_upload_media(self, upload_file, qualified_url, validate, fallback_url):
-        if self["supportsMediaUpload"] is not True:
-            if validate is True:
+        if not self["supportsMediaUpload"]:
+            if validate:
                 raise ValidationError(
                     "upload_file was provided while method doesn't support media upload"
                 )
-            else:
-                # This will probably not work, but will return a mediaupload object anyway
-                return MediaUpload(upload_file, upload_path=fallback_url)
+            # This will probably not work, but will return a mediaupload object anyway
+            return MediaUpload(upload_file, upload_path=fallback_url)
 
         # If resumable, create resumable object
         resumable = (
-            self._build_resumeable_media(upload_file, qualified_url)
+            self._build_resumeable_media(qualified_url)
             if _safe_getitem(
                 self._method_specs, "mediaUpload", "protocols", "resumable"
             )
@@ -631,8 +630,6 @@ class Method:
         )
 
         # Create MediaUpload object and pass it the resumable object we just created
-        # simple_upload_path = _safe_getitem(self._method_specs, 'mediaUpload', 'protocols', 'simple', 'path')
-        # media_upload_url_base = self._root_url[:-1] + simple_upload_path
         media_upload_url_base = self._root_url + "upload/" + self._service_path
         media_upload_url = qualified_url.replace(self._base_url, media_upload_url_base)
         max_size = self._convert_str_size_to_int(
@@ -642,7 +639,7 @@ class Method:
         multipart = self["mediaUpload"]["protocols"]["simple"].get("multipart", True)
 
         return MediaUpload(
-            file_path=upload_file,
+            upload_file,
             upload_path=media_upload_url,
             max_size=max_size,
             mime_range=mime_range,
@@ -651,15 +648,11 @@ class Method:
             validate=validate,
         )
 
-    def _build_resumeable_media(self, upload_file, qualified_url):
-        # resumable_upload_path = _safe_getitem(self._method_specs, 'mediaUpload', 'protocols', 'resumable', 'path')
-        # resumable_url_base = self._root_url[:-1] + resumable_upload_path
+    def _build_resumeable_media(self, qualified_url):
         resumable_url_base = self._root_url + "resumable/upload/" + self._service_path
         resumable_url = qualified_url.replace(self._base_url, resumable_url_base)
         multipart = self["mediaUpload"]["protocols"]["resumable"].get("multipart", True)
-        return ResumableUpload(
-            upload_file, multipart=multipart, upload_path=resumable_url
-        )
+        return ResumableUpload(multipart=multipart, upload_path=resumable_url)
 
     @staticmethod
     def _convert_str_size_to_int(size):
diff --git a/aiogoogle/sessions/aiohttp_session.py b/aiogoogle/sessions/aiohttp_session.py
index d189d38..4706658 100644
--- a/aiogoogle/sessions/aiohttp_session.py
+++ b/aiogoogle/sessions/aiohttp_session.py
@@ -11,7 +11,6 @@ import async_timeout
 
 from ..models import Response
 from .abc import AbstractSession
-from ..excs import ValidationError
 
 
 async def _get_file_size(full_file_path):
@@ -28,6 +27,11 @@ async def _aiter_file(file_name, chunk_size):
             chunk = await f.read(chunk_size)
 
 
+async def _read_file(file_name):
+    async with aiofiles.open(file_name, "rb") as file:
+        return await file.read()
+
+
 class AiohttpSession(ClientSession, AbstractSession):
     async def send(
         self,
@@ -84,16 +88,7 @@ class AiohttpSession(ClientSession, AbstractSession):
             request.headers["User-Agent"] = "Aiogoogle Aiohttp (gzip)"
             if request.media_upload:
                 # Validate
-                if (
-                    request.media_upload.validate is True
-                    and request.media_upload.max_size is not None
-                ):
-                    size = await _get_file_size(request.media_upload.file_path)
-                    max_size = request.media_upload.max_size
-                    if size > max_size:
-                        raise ValidationError(
-                            f'"{request.media_upload.file_path}" has a size of {size/1000}KB. Max upload size for this endpoint is: {max_size/1000}KB.'
-                        )
+                await request.media_upload.run_validation(_get_file_size)
 
                 # If multipart pass a file async generator
                 if request.media_upload.multipart is True:
@@ -101,10 +96,7 @@ class AiohttpSession(ClientSession, AbstractSession):
                         mpwriter.append_json(request.json)
 
                         mpwriter.append(
-                            _aiter_file(
-                                request.media_upload.file_path,
-                                request.media_upload.chunk_size
-                            ),
+                            request.media_upload.aiter_file(_aiter_file),
                             headers={"Content-Type": request.upload_file_content_type} if request.upload_file_content_type else None
                         )
 
@@ -122,21 +114,18 @@ class AiohttpSession(ClientSession, AbstractSession):
                         )
                 # Else load file to memory and send
                 else:
-                    async with aiofiles.open(
-                        request.media_upload.file_path, "rb"
-                    ) as file:
-                        read_file = await file.read()
-                        if request.upload_file_content_type:
-                            request.headers.update({"Content-Type": request.upload_file_content_type})
-                        return await self.request(
-                            method=request.method,
-                            url=request.media_upload.upload_path,
-                            headers=request.headers,
-                            data=read_file,
-                            json=request.json,
-                            timeout=request.timeout,
-                            verify_ssl=request._verify_ssl,
-                        )
+                    read_file = await request.media_upload.read_file(_read_file)
+                    if request.upload_file_content_type:
+                        request.headers.update({"Content-Type": request.upload_file_content_type})
+                    return await self.request(
+                        method=request.method,
+                        url=request.media_upload.upload_path,
+                        headers=request.headers,
+                        data=read_file,
+                        json=request.json,
+                        timeout=request.timeout,
+                        verify_ssl=request._verify_ssl,
+                    )
             # Else, if no file upload
             else:
                 return await self.request(


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index f81f795..aaa1213 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,5 +1,5 @@
-import os
 import json
+from pathlib import Path
 
 import pytest
 
@@ -9,13 +9,9 @@ from aiogoogle.resource import GoogleAPI
 @pytest.fixture
 def open_discovery_document():
     def wrapped(name, version):
-        current_dir = os.getcwd()
-        file_name = (
-            current_dir + "/tests/data/" + name + "_" + version + "_discovery_doc.json"
-        )
-        with open(file_name, "r") as discovery_doc:
-            discovery_doc_dict = json.loads(discovery_doc.read())
-        return discovery_doc_dict
+        json_file = Path(__file__).parent / "data" / f"{name}_{version}_discovery_doc.json"
+        with json_file.open("r") as discovery_doc:
+            return json.loads(discovery_doc.read())
 
     return wrapped
 
diff --git a/tests/refresh_all_apis.py b/tests/refresh_all_apis.py
index 8499c2a..21b99f6 100755
--- a/tests/refresh_all_apis.py
+++ b/tests/refresh_all_apis.py
@@ -1,41 +1,25 @@
 #!/usr/bin/python3.7
-
-from aiogoogle import Aiogoogle
-import os
-import sys
-import errno
+from pathlib import Path
 import json
 import asyncio
-from aiohttp import ClientSession
-from aiogoogle import HTTPError
 import pprint
 
-
-def _check_for_correct_cwd(current_dir):
-    if current_dir[-9:] != "aiogoogle":  # current dir is aiogoogle
-        print(current_dir)
-        print("must be in aiogoogle's dir, not test dir")
-        sys.exit()
+from aiohttp import ClientSession
+from aiogoogle import Aiogoogle
+from aiogoogle import HTTPError
 
 
 def _pop_unstable_apis(all_apis: list):
-    stable_apis = []
-    for api in all_apis:
-        if not len(api[1]) > 3:  # No funky versions because they break the tests alot
-            stable_apis.append(api)
-    return stable_apis
+    return [api for api in all_apis if len(api[1]) <= 3]
 
 
 async def refresh_disc_docs_json():
     file_errors = []
-    current_dir = os.getcwd()
 
     # Create new .data/ dir if one doesn't exists
-    _check_for_correct_cwd(current_dir)
-
-    # Refresh all_apis in tests/tests_globals.py
-    all_apis = []
-    final_all_apis = []
+    this_folder = Path(__file__).parent
+    data_dir = this_folder / "data"
+    data_dir.mkdir(exist_ok=True)
 
     async with ClientSession() as sess:
         apis_pref = await sess.get(
@@ -43,9 +27,8 @@ async def refresh_disc_docs_json():
         )
         apis_pref = await apis_pref.json()
 
-    for api in apis_pref["items"]:
-        all_apis.append((api["name"], api["version"]))
-
+    # Refresh all_apis in tests/tests_globals.py
+    all_apis = [(api["name"], api["version"]) for api in apis_pref["items"]]
     all_apis = _pop_unstable_apis(all_apis)
     final_all_apis = all_apis
 
@@ -55,16 +38,13 @@ async def refresh_disc_docs_json():
         all_discovery_documents = await asyncio.gather(*tasks, return_exceptions=True)
 
     # Refresh discovery files in tests/data
-    for i, google_api in enumerate(all_discovery_documents):
-        name = all_apis[i][0]
-        version = all_apis[i][1]
+    for google_api, (name, version) in zip(all_discovery_documents, all_apis):
         if isinstance(google_api, HTTPError):
             e = google_api
             if e.res.status_code != 404:
                 print('Non 404 error')
                 print('\033[91m\n' + e + '\n\033[0m')
-
-            if e.res.status_code == 404:
+            else:
                 # only ignore if it's a 404 error. Should raise an error otherwise
                 final_all_apis = list(filter(lambda api: (api[0] != name), final_all_apis))
 
@@ -72,32 +52,17 @@ async def refresh_disc_docs_json():
             print(f'\033[91mError: Failed to download {name} {version}\033[0m')
             continue
 
-        data_dir_name = current_dir + "/tests/data/"
-        try:
-            if not os.path.exists(data_dir_name):
-                os.makedirs(data_dir_name)
-        except OSError as e:
-            if e.errno != errno.EEXIST:
-                raise
-
         # Save discovery docuemnt as .json file to the newly created data dir
-        file_name = (
-            current_dir
-            + "/tests/data/"
-            + name
-            + "_"
-            + version
-            + "_discovery_doc.json"
-        )
-        with open(file_name, "w") as discovery_file:
+        json_file = data_dir / f"{name}_{version}_discovery_doc.json"
+        with open(str(json_file), "w") as discovery_file:
             json.dump(google_api.discovery_document, discovery_file)
-        print(f"saved {name}-{version} to {file_name}")
+        print(f"saved {name}-{version} to {json_file}")
 
     print("Done")
     if file_errors:
         print(f"Errors found: {str(file_errors)}")
 
-    with open("tests/ALL_APIS.py", "w") as f:
+    with this_folder.joinpath("ALL_APIS.py").open("w") as f:
         f.write("""### This file is autogenerated ###\n""")
         f.write(f"ALL_APIS = {pprint.pformat(final_all_apis)}")
         print("SUCCESS!")
diff --git a/tests/test_units/test_models.py b/tests/test_units/test_models.py
index 5739cfd..0d7dc25 100644
--- a/tests/test_units/test_models.py
+++ b/tests/test_units/test_models.py
@@ -4,28 +4,108 @@
 # TODO: Test next page generator
 # TODO: Test response.__call__ returns pagination gen
 # TODO: Test next_page returns a valid request
-from aiogoogle.models import Request
+import asyncio
+from unittest.mock import MagicMock
 
+import pytest
 
-def test_request_add_query_param():
-    r = Request(url="https://example.com")
-    r._add_query_param({'foo': 'bar'})
-    assert r.url == 'https://example.com?foo=bar'
+from aiogoogle.models import Request, MediaUpload
+from aiogoogle.excs import ValidationError
 
 
-def test_request_add_query_param1():
-    r = Request(url="https://example.com/foo")
-    r._add_query_param({'bar': 'baz'})
-    assert r.url == 'https://example.com/foo?bar=baz'
+def async_return(result):
+    f = asyncio.Future()
+    f.set_result(result)
+    return f
 
 
-def test_request_add_query_param2():
-    r = Request(url="https://example.com/")
-    r._add_query_param({'bar': 'baz'})
-    assert r.url == 'https://example.com?bar=baz'
+class AsyncIterator:
+    def __init__(self, seq):
+        self.iter = iter(seq)
 
+    def __aiter__(self):
+        return self
 
-def test_request_add_query_param3():
-    r = Request(url="https://example.com/foo?bar=baz")
-    r._add_query_param({'idk': 'fuuu'})
-    assert r.url == 'https://example.com/foo?bar=baz&idk=fuuu'
+    async def __anext__(self):
+        try:
+            return next(self.iter)
+        except StopIteration:
+            raise StopAsyncIteration
+
+
+add_query_params = (
+    ("https://example.com", {'foo': 'bar'}, 'https://example.com?foo=bar'),
+    ("https://example.com/foo", {'bar': 'baz'}, 'https://example.com/foo?bar=baz'),
+    ("https://example.com/", {'bar': 'baz'}, 'https://example.com?bar=baz'),
+    ("https://example.com/foo?bar=baz", {'idk': 'fuuu'}, 'https://example.com/foo?bar=baz&idk=fuuu')
+)
+
+
+@pytest.mark.parametrize("url, params, result", add_query_params)
+def test_request_add_query_param(url, params, result):
+    r = Request(url=url)
+    r._add_query_param(params)
+    assert r.url == result
+
+
+def noop(*args, **kwargs):
+    return
+
+
+@pytest.mark.asyncio
+async def test_media_upload_bytes_aiter():
+    body = bytes(101)
+    mu = MediaUpload(body, chunk_size=10)
+    chunks = [chunk async for chunk in mu.aiter_file(noop)]
+    assert len(chunks) == 11
+    joined = b"".join(chunks)
+    assert body == joined
+
+
+@pytest.mark.asyncio
+async def test_media_upload_bytes_read():
+    body = bytes(101)
+    mu = MediaUpload(body)
+    assert await mu.read_file(noop) == body
+
+
+@pytest.mark.asyncio
+async def test_media_upload_bytes_validation():
+    body = bytes(101)
+    mu = MediaUpload(body, max_size=10)
+    with pytest.raises(ValidationError):
+        await mu.run_validation(noop)
+
+    mu = MediaUpload(body, max_size=200)
+    await mu.run_validation(noop)
+
+
+@pytest.mark.asyncio
+async def test_media_upload_aiter():
+    file_path = "test"
+    mu = MediaUpload(file_path, chunk_size=10)
+    aiter_func = MagicMock(return_value=AsyncIterator(range(5)))
+    chunks = [chunk async for chunk in mu.aiter_file(aiter_func)]
+    aiter_func.assert_called_once_with(file_path, 10)
+    assert len(chunks) == 5
+
+
+@pytest.mark.asyncio
+async def test_media_upload_read():
+    file_path = "test"
+    mu = MediaUpload(file_path)
+    read_func = MagicMock(return_value=async_return(None))
+    await mu.read_file(read_func)
+    read_func.assert_called_once_with(file_path)
+
+
+@pytest.mark.asyncio
+async def test_media_upload_validation():
+    file_path = "test"
+    mu = MediaUpload(file_path, max_size=10)
+    size_func = MagicMock(return_value=async_return(100))
+    with pytest.raises(ValidationError):
+        await mu.run_validation(size_func)
+
+    mu = MediaUpload(file_path, max_size=200)
+    await mu.run_validation(size_func)