Merge pull request #93 from NabilMostafa/stream-uploaded-files

Stream uploaded files

aiogoogle/models.py

@@ -52,6 +52,7 @@ class MediaUpload:
 
         validate (bool): Whether or not a session should validate the upload size before sending
 
+        pipe_from (file object): class object to stream file content from
     """
 
     def __init__(
@@ -64,6 +65,7 @@ def __init__(
         chunk_size=None,
         resumable=None,
         validate=True,
+        pipe_from=None
     ):
         if isinstance(file_path_or_bytes, bytes):
             self.file_body = file_path_or_bytes
@@ -78,6 +80,7 @@ def __init__(
         self.chunk_size = chunk_size or DEFAULT_UPLOAD_CHUNK_SIZE
         self.resumable = resumable
         self.validate = validate
+        self.pipe_from = pipe_from
 
     async def run_validation(self, size_func):
         if self.validate and self.max_size:
@@ -93,6 +96,8 @@ 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
+        elif self.pipe_from:
+            yield self.pipe_from.read()
         else:
             async for chunk in self._aiter_body():
                 yield chunk
@@ -265,6 +270,8 @@ class Response:
 
         upload_file (str): path of the upload file specified in the request
 
+        pipe_from (file object): class object to stream file content from
+
         session_factory (aiogoogle.sessions.abc.AbstractSession): A callable implementation of aiogoogle's session interface
     """
 
@@ -280,6 +287,7 @@ def __init__(
         download_file=None,
         pipe_to=None,
         upload_file=None,
+        pipe_from=None,
         session_factory=None,
     ):
         if json and data:
@@ -295,6 +303,7 @@ def __init__(
         self.download_file = download_file
         self.pipe_to = pipe_to
         self.upload_file = upload_file
+        self.pipe_from = pipe_from
         self.session_factory = session_factory
 
     @staticmethod

aiogoogle/resource.py

@@ -22,6 +22,7 @@
     "data",
     "json",
     "upload_file",
+    "pipe_from",
     "download_file",
     "pipe_to",
     "timeout",
@@ -60,6 +61,7 @@ def wrapper(
         data=None,
         json=None,
         upload_file=None,
+        pipe_from=None,
         download_file=None,
         pipe_to=None,
         timeout=None,
@@ -85,6 +87,7 @@ def wrapper(
             data,
             json,
             upload_file,
+            pipe_from,
             download_file,
             pipe_to,
             timeout,
@@ -389,6 +392,7 @@ def __call__(
         data=None,
         json=None,
         upload_file=None,
+        pipe_from=None,
         download_file=None,
         pipe_to=None,
         timeout=None,
@@ -579,13 +583,20 @@ def __call__(
                     )
             media_download = MediaDownload(pipe_to=pipe_to)
 
+        if not upload_file and not pipe_from:
+            media_upload = None
+
         # Process upload_file
         if upload_file:
             media_upload = self._build_upload_media(
-                upload_file, uri, validate, fallback_url=url
+                upload_file, uri, validate, fallback_url=url, pipe_from=None
+            )
+
+        # Process pipe_from
+        if pipe_from:
+            media_upload = self._build_upload_media(
+                upload_file, uri, validate, fallback_url=url, pipe_from=pipe_from
             )
-        else:
-            media_upload = None
 
         return Request(
             method=self["httpMethod"],
@@ -631,7 +642,7 @@ def _build_url(self, base_url, uri_params, validate):
         else:
             return base_url + self["path"]
 
-    def _build_upload_media(self, upload_file, qualified_url, validate, fallback_url):
+    def _build_upload_media(self, upload_file, qualified_url, validate, fallback_url, pipe_from):
         if not self["supportsMediaUpload"]:
             if validate:
                 raise ValidationError(
@@ -660,6 +671,7 @@ def _build_upload_media(self, upload_file, qualified_url, validate, fallback_url
 
         return MediaUpload(
             upload_file,
+            pipe_from=pipe_from,
             upload_path=media_upload_url,
             max_size=max_size,
             mime_range=mime_range,

aiogoogle/sessions/aiohttp_session.py

@@ -46,6 +46,7 @@ async def resolve_response(request, response):
             json = None
             download_file = None
             upload_file = None
+            pipe_from = None
             pipe_to = None
 
             # If downloading file:
@@ -89,6 +90,7 @@ async def resolve_response(request, response):
                 download_file=download_file,
                 pipe_to=pipe_to,
                 upload_file=upload_file,
+                pipe_from=pipe_from,
                 session_factory=session_factory,
             )
 

aiogoogle/sessions/curio_asks_session.py

@@ -30,6 +30,7 @@ async def resolve_response(request, response):
             json = None
             download_file = None
             upload_file = None
+            pipe_from = None
 
             # If downloading file:
             if request.media_download:
@@ -67,6 +68,7 @@ async def resolve_response(request, response):
                 req=request,
                 download_file=download_file,
                 upload_file=upload_file,
+                pipe_from=pipe_from,
                 session_factory=session_factory,
             )
 

examples/stream_upload_file.py

@@ -0,0 +1,40 @@
+#!/usr/bin/python3.7
+
+'''
+Scopes Required:
+
+* https://www.googleapis.com/auth/drive
+* https://www.googleapis.com/auth/drive.file
+
+API explorer link:
+
+* https://developers.google.com/apis-explorer/#p/drive/v3/drive.files.create
+'''
+
+import asyncio
+
+from helpers import Aiogoogle, user_creds, client_creds
+
+
+class MyFile:
+    @staticmethod
+    async def read():
+        return b'Hello World'
+
+
+async def stream_upload_file():
+    async with Aiogoogle(user_creds=user_creds, client_creds=client_creds) as aiogoogle:
+        # Create API
+        drive_v3 = await aiogoogle.discover("drive", "v3")
+
+        req = drive_v3.files.create(
+            pipe_from=MyFile(),
+            fields="id"
+        )
+
+        # Upload file
+        upload_res = await aiogoogle.as_user(req)
+        print(f"Uploaded stream file successfully.\nFile ID: {upload_res['id']}")
+
+if __name__ == "__main__":
+    asyncio.run(stream_upload_file())