Clean up the file upload interface with FileVar class (#549)

Author: leszekhanusz

docs/usage/file_upload.rst

@@ -14,11 +14,14 @@ Single File
 In order to upload a single file, you need to:
 
 * set the file as a variable value in the mutation
-* provide the opened file to the `variable_values` argument of `execute`
+* create a :class:`FileVar <gql.FileVar>` object with your file path
+* provide the `FileVar` instance to the `variable_values` argument of `execute`
 * set the `upload_files` argument to True
 
 .. code-block:: python
 
+    from gql import client, gql, FileVar
+
     transport = AIOHTTPTransport(url='YOUR_URL')
     # Or transport = RequestsHTTPTransport(url='YOUR_URL')
     # Or transport = HTTPXTransport(url='YOUR_URL')
@@ -34,32 +37,38 @@ In order to upload a single file, you need to:
       }
     ''')
 
-    with open("YOUR_FILE_PATH", "rb") as f:
-
-        params = {"file": f}
+    params = {"file": FileVar("YOUR_FILE_PATH")}
 
-        result = client.execute(
-            query, variable_values=params, upload_files=True
-        )
+    result = client.execute(
+        query, variable_values=params, upload_files=True
+    )
 
 Setting the content-type
 ^^^^^^^^^^^^^^^^^^^^^^^^
 
 If you need to set a specific Content-Type attribute to a file,
-you can set the :code:`content_type` attribute of the file like this:
+you can set the :code:`content_type` attribute of :class:`FileVar <gql.FileVar>`:
 
 .. code-block:: python
 
-    with open("YOUR_FILE_PATH", "rb") as f:
+    # Setting the content-type to a pdf file for example
+    filevar = FileVar(
+        "YOUR_FILE_PATH",
+        content_type="application/pdf",
+    )
 
-        # Setting the content-type to a pdf file for example
-        f.content_type = "application/pdf"
+Setting the uploaded file name
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-        params = {"file": f}
+To modify the uploaded filename, use the :code:`filename` attribute of :class:`FileVar <gql.FileVar>`:
 
-        result = client.execute(
-            query, variable_values=params, upload_files=True
-        )
+.. code-block:: python
+
+    # Setting the content-type to a pdf file for example
+    filevar = FileVar(
+        "YOUR_FILE_PATH",
+        filename="filename1.txt",
+    )
 
 File list
 ---------
@@ -68,6 +77,8 @@ It is also possible to upload multiple files using a list.
 
 .. code-block:: python
 
+    from gql import client, gql, FileVar
+
     transport = AIOHTTPTransport(url='YOUR_URL')
     # Or transport = RequestsHTTPTransport(url='YOUR_URL')
     # Or transport = HTTPXTransport(url='YOUR_URL')
@@ -83,18 +94,15 @@ It is also possible to upload multiple files using a list.
       }
     ''')
 
-    f1 = open("YOUR_FILE_PATH_1", "rb")
-    f2 = open("YOUR_FILE_PATH_2", "rb")
+    f1 = FileVar("YOUR_FILE_PATH_1")
+    f2 = FileVar("YOUR_FILE_PATH_2")
 
     params = {"files": [f1, f2]}
 
     result = client.execute(
         query, variable_values=params, upload_files=True
     )
 
-    f1.close()
-    f2.close()
-
 
 Streaming
 ---------
@@ -120,18 +128,8 @@ Streaming local files
 aiohttp allows to upload files using an asynchronous generator.
 See `Streaming uploads on aiohttp docs`_.
 
-
-In order to stream local files, instead of providing opened files to the
-`variable_values` argument of `execute`, you need to provide an async generator
-which will provide parts of the files.
-
-You can use `aiofiles`_
-to read the files in chunks and create this asynchronous generator.
-
-.. _Streaming uploads on aiohttp docs: https://docs.aiohttp.org/en/stable/client_quickstart.html#streaming-uploads
-.. _aiofiles: https://github.com/Tinche/aiofiles
-
-Example:
+From gql version 4.0, it is possible to activate file streaming simply by
+setting the `streaming` argument of :class:`FileVar <gql.FileVar>` to `True`
 
 .. code-block:: python
 
@@ -147,18 +145,38 @@ Example:
       }
     ''')
 
+    f1 = FileVar(
+        file_name='YOUR_FILE_PATH',
+        streaming=True,
+    )
+
+    params = {"file": f1}
+
+    result = client.execute(
+        query, variable_values=params, upload_files=True
+    )
+
+Another option is to use an async generator to provide parts of the file.
+
+You can use `aiofiles`_
+to read the files in chunks and create this asynchronous generator.
+
+.. _Streaming uploads on aiohttp docs: https://docs.aiohttp.org/en/stable/client_quickstart.html#streaming-uploads
+.. _aiofiles: https://github.com/Tinche/aiofiles
+
+.. code-block:: python
+
     async def file_sender(file_name):
         async with aiofiles.open(file_name, 'rb') as f:
-            chunk = await f.read(64*1024)
-                while chunk:
-                    yield chunk
-                    chunk = await f.read(64*1024)
+            while chunk := await f.read(64*1024):
+                yield chunk
 
-    params = {"file": file_sender(file_name='YOUR_FILE_PATH')}
+    f1 = FileVar(file_sender(file_name='YOUR_FILE_PATH'))
+    params = {"file": f1}
 
     result = client.execute(
-		query, variable_values=params, upload_files=True
-	)
+        query, variable_values=params, upload_files=True
+    )
 
 Streaming downloaded files
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -200,7 +218,7 @@ Example:
               }
             ''')
 
-            params = {"file": resp.content}
+            params = {"file": FileVar(resp.content)}
 
             result = client.execute(
                 query, variable_values=params, upload_files=True

gql/__init__.py

@@ -11,10 +11,12 @@
 from .client import Client
 from .gql import gql
 from .graphql_request import GraphQLRequest
+from .transport.file_upload import FileVar
 
 __all__ = [
     "__version__",
     "gql",
     "Client",
     "GraphQLRequest",
+    "FileVar",
 ]

gql/transport/aiohttp.py

@@ -23,7 +23,6 @@
 from graphql import DocumentNode, ExecutionResult, print_ast
 from multidict import CIMultiDictProxy
 
-from ..utils import extract_files
 from .appsync_auth import AppSyncAuthentication
 from .async_transport import AsyncTransport
 from .common.aiohttp_closed_event import create_aiohttp_closed_event
@@ -33,6 +32,7 @@
     TransportProtocolError,
     TransportServerError,
 )
+from .file_upload import FileVar, close_files, extract_files, open_files
 
 log = logging.getLogger(__name__)
 
@@ -207,6 +207,10 @@ async def execute(
                 file_classes=self.file_classes,
             )
 
+            # Opening the files using the FileVar parameters
+            open_files(list(files.values()), transport_supports_streaming=True)
+            self.files = files
+
             # Save the nulled variable values in the payload
             payload["variables"] = nulled_variable_values
 
@@ -220,8 +224,8 @@ async def execute(
             file_map = {str(i): [path] for i, path in enumerate(files)}
 
             # Enumerate the file streams
-            # Will generate something like {'0': <_io.BufferedReader ...>}
-            file_streams = {str(i): files[path] for i, path in enumerate(files)}
+            # Will generate something like {'0': FileVar object}
+            file_vars = {str(i): files[path] for i, path in enumerate(files)}
 
             # Add the payload to the operations field
             operations_str = self.json_serialize(payload)
@@ -235,12 +239,15 @@ async def execute(
             log.debug("file_map %s", file_map_str)
             data.add_field("map", file_map_str, content_type="application/json")
 
-            # Add the extracted files as remaining fields
-            for k, f in file_streams.items():
-                name = getattr(f, "name", k)
-                content_type = getattr(f, "content_type", None)
+            for k, file_var in file_vars.items():
+                assert isinstance(file_var, FileVar)
 
-                data.add_field(k, f, filename=name, content_type=content_type)
+                data.add_field(
+                    k,
+                    file_var.f,
+                    filename=file_var.filename,
+                    content_type=file_var.content_type,
+                )
 
             post_args: Dict[str, Any] = {"data": data}
 
@@ -267,51 +274,59 @@ async def execute(
         if self.session is None:
             raise TransportClosed("Transport is not connected")
 
-        async with self.session.post(self.url, ssl=self.ssl, **post_args) as resp:
-
-            # Saving latest response headers in the transport
-            self.response_headers = resp.headers
+        try:
+            async with self.session.post(self.url, ssl=self.ssl, **post_args) as resp:
 
-            async def raise_response_error(
-                resp: aiohttp.ClientResponse, reason: str
-            ) -> NoReturn:
-                # We raise a TransportServerError if the status code is 400 or higher
-                # We raise a TransportProtocolError in the other cases
+                # Saving latest response headers in the transport
+                self.response_headers = resp.headers
 
-                try:
-                    # Raise a ClientResponseError if response status is 400 or higher
-                    resp.raise_for_status()
-                except ClientResponseError as e:
-                    raise TransportServerError(str(e), e.status) from e
-
-                result_text = await resp.text()
-                raise TransportProtocolError(
-                    f"Server did not return a GraphQL result: "
-                    f"{reason}: "
-                    f"{result_text}"
-                )
+                async def raise_response_error(
+                    resp: aiohttp.ClientResponse, reason: str
+                ) -> NoReturn:
+                    # We raise a TransportServerError if status code is 400 or higher
+                    # We raise a TransportProtocolError in the other cases
 
-            try:
-                result = await resp.json(loads=self.json_deserialize, content_type=None)
+                    try:
+                        # Raise ClientResponseError if response status is 400 or higher
+                        resp.raise_for_status()
+                    except ClientResponseError as e:
+                        raise TransportServerError(str(e), e.status) from e
 
-                if log.isEnabledFor(logging.INFO):
                     result_text = await resp.text()
-                    log.info("<<< %s", result_text)
+                    raise TransportProtocolError(
+                        f"Server did not return a GraphQL result: "
+                        f"{reason}: "
+                        f"{result_text}"
+                    )
 
-            except Exception:
-                await raise_response_error(resp, "Not a JSON answer")
+                try:
+                    result = await resp.json(
+                        loads=self.json_deserialize, content_type=None
+                    )
 
-            if result is None:
-                await raise_response_error(resp, "Not a JSON answer")
+                    if log.isEnabledFor(logging.INFO):
+                        result_text = await resp.text()
+                        log.info("<<< %s", result_text)
 
-            if "errors" not in result and "data" not in result:
-                await raise_response_error(resp, 'No "data" or "errors" keys in answer')
+                except Exception:
+                    await raise_response_error(resp, "Not a JSON answer")
 
-            return ExecutionResult(
-                errors=result.get("errors"),
-                data=result.get("data"),
-                extensions=result.get("extensions"),
-            )
+                if result is None:
+                    await raise_response_error(resp, "Not a JSON answer")
+
+                if "errors" not in result and "data" not in result:
+                    await raise_response_error(
+                        resp, 'No "data" or "errors" keys in answer'
+                    )
+
+                return ExecutionResult(
+                    errors=result.get("errors"),
+                    data=result.get("data"),
+                    extensions=result.get("extensions"),
+                )
+        finally:
+            if upload_files:
+                close_files(list(self.files.values()))
 
     def subscribe(
         self,