Removing _prepare_batch_request to avoid code duplication

Author: leszekhanusz

gql/transport/aiohttp.py

@@ -127,7 +127,7 @@ async def connect(self) -> None:
 
             # Adding custom parameters passed from init
             if self.client_session_args:
-                client_session_args.update(self.client_session_args)  # type: ignore
+                client_session_args.update(self.client_session_args)
 
             log.debug("Connecting transport")
 
@@ -164,36 +164,22 @@ async def close(self) -> None:
 
         self.session = None
 
-    def _prepare_batch_request(
-        self,
-        reqs: List[GraphQLRequest],
-        extra_args: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, Any]:
-
-        payload = [req.payload for req in reqs]
-
-        post_args = {"json": payload}
-
-        # Log the payload
-        if log.isEnabledFor(logging.DEBUG):
-            log.debug(">>> %s", self.json_serialize(payload))
-
-        # Pass post_args to aiohttp post method
-        if extra_args:
-            post_args.update(extra_args)
-
-        return post_args
-
     def _prepare_request(
         self,
-        request: GraphQLRequest,
+        request: Union[GraphQLRequest, List[GraphQLRequest]],
         extra_args: Optional[Dict[str, Any]] = None,
         upload_files: bool = False,
     ) -> Dict[str, Any]:
 
-        payload = request.payload
+        payload: Dict | List
+        if isinstance(request, GraphQLRequest):
+            payload = request.payload
+        else:
+            payload = [req.payload for req in request]
 
         if upload_files:
+            assert isinstance(payload, Dict)
+            assert isinstance(request, GraphQLRequest)
             post_args = self._prepare_file_uploads(request, payload)
         else:
             post_args = {"json": payload}
@@ -416,7 +402,7 @@ async def execute_batch(
         if self.session is None:
             raise TransportClosed("Transport is not connected")
 
-        post_args = self._prepare_batch_request(
+        post_args = self._prepare_request(
             reqs,
             extra_args,
         )

gql/transport/httpx.py

@@ -59,15 +59,22 @@ def __init__(
 
     def _prepare_request(
         self,
-        req: GraphQLRequest,
+        request: Union[GraphQLRequest, List[GraphQLRequest]],
+        *,
         extra_args: Optional[Dict[str, Any]] = None,
         upload_files: bool = False,
     ) -> Dict[str, Any]:
 
-        payload = req.payload
+        payload: Dict | List
+        if isinstance(request, GraphQLRequest):
+            payload = request.payload
+        else:
+            payload = [req.payload for req in request]
 
         if upload_files:
-            post_args = self._prepare_file_uploads(req, payload)
+            assert isinstance(payload, Dict)
+            assert isinstance(request, GraphQLRequest)
+            post_args = self._prepare_file_uploads(request, payload)
         else:
             post_args = {"json": payload}
 
@@ -81,26 +88,6 @@ def _prepare_request(
 
         return post_args
 
-    def _prepare_batch_request(
-        self,
-        reqs: List[GraphQLRequest],
-        extra_args: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, Any]:
-
-        payload = [req.payload for req in reqs]
-
-        post_args = {"json": payload}
-
-        # Log the payload
-        if log.isEnabledFor(logging.DEBUG):
-            log.debug(">>> %s", self.json_serialize(payload))
-
-        # Pass post_args to httpx post method
-        if extra_args:
-            post_args.update(extra_args)
-
-        return post_args
-
     def _prepare_file_uploads(
         self,
         request: GraphQLRequest,
@@ -244,7 +231,7 @@ def connect(self):
 
         self.client = httpx.Client(**self.kwargs)
 
-    def execute(  # type: ignore
+    def execute(
         self,
         request: GraphQLRequest,
         *,
@@ -269,8 +256,8 @@ def execute(  # type: ignore
 
         post_args = self._prepare_request(
             request,
-            extra_args,
-            upload_files,
+            extra_args=extra_args,
+            upload_files=upload_files,
         )
 
         try:
@@ -302,9 +289,9 @@ def execute_batch(
         if not self.client:
             raise TransportClosed("Transport is not connected")
 
-        post_args = self._prepare_batch_request(
+        post_args = self._prepare_request(
             reqs,
-            extra_args,
+            extra_args=extra_args,
         )
 
         response = self.client.post(self.url, **post_args)
@@ -361,8 +348,8 @@ async def execute(
 
         post_args = self._prepare_request(
             request,
-            extra_args,
-            upload_files,
+            extra_args=extra_args,
+            upload_files=upload_files,
         )
 
         try:
@@ -394,9 +381,9 @@ async def execute_batch(
         if not self.client:
             raise TransportClosed("Transport is not connected")
 
-        post_args = self._prepare_batch_request(
+        post_args = self._prepare_request(
             reqs,
-            extra_args,
+            extra_args=extra_args,
         )
 
         response = await self.client.post(self.url, **post_args)

gql/transport/requests.py

@@ -137,50 +137,20 @@ def connect(self):
         else:
             raise TransportAlreadyConnected("Transport is already connected")
 
-    def _prepare_batch_request(
-        self,
-        reqs: List[GraphQLRequest],
-        *,
-        timeout: Optional[int] = None,
-        extra_args: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, Any]:
-
-        payload = [req.payload for req in reqs]
-
-        post_args: Dict[str, Any] = {
-            "headers": self.headers,
-            "auth": self.auth,
-            "cookies": self.cookies,
-            "timeout": timeout or self.default_timeout,
-            "verify": self.verify,
-        }
-
-        data_key = "json" if self.use_json else "data"
-        post_args[data_key] = payload
-
-        # Log the payload
-        if log.isEnabledFor(logging.DEBUG):
-            log.debug(">>> %s", self.json_serialize(payload))
-
-        # Pass kwargs to requests post method
-        post_args.update(self.kwargs)
-
-        # Pass post_args to requests post method
-        if extra_args:
-            post_args.update(extra_args)
-
-        return post_args
-
     def _prepare_request(
         self,
-        request: GraphQLRequest,
+        request: Union[GraphQLRequest, List[GraphQLRequest]],
         *,
         timeout: Optional[int] = None,
         extra_args: Optional[Dict[str, Any]] = None,
         upload_files: bool = False,
     ) -> Dict[str, Any]:
 
-        payload = request.payload
+        payload: Dict | List
+        if isinstance(request, GraphQLRequest):
+            payload = request.payload
+        else:
+            payload = [req.payload for req in request]
 
         post_args: Dict[str, Any] = {
             "headers": self.headers,
@@ -191,6 +161,8 @@ def _prepare_request(
         }
 
         if upload_files:
+            assert isinstance(payload, Dict)
+            assert isinstance(request, GraphQLRequest)
             post_args = self._prepare_file_uploads(
                 request=request,
                 payload=payload,
@@ -282,7 +254,7 @@ def _prepare_file_uploads(
 
         return post_args
 
-    def execute(  # type: ignore
+    def execute(
         self,
         request: GraphQLRequest,
         timeout: Optional[int] = None,
@@ -316,9 +288,7 @@ def execute(  # type: ignore
 
         # Using the created session to perform requests
         try:
-            response = self.session.request(
-                self.method, self.url, **post_args  # type: ignore
-            )
+            response = self.session.request(self.method, self.url, **post_args)
         finally:
             if upload_files:
                 close_files(list(self.files.values()))
@@ -373,7 +343,7 @@ def execute_batch(
         if not self.session:
             raise TransportClosed("Transport is not connected")
 
-        post_args = self._prepare_batch_request(
+        post_args = self._prepare_request(
             reqs,
             timeout=timeout,
             extra_args=extra_args,