Finish type checking

Author: HungKNguyen

src/nutrient_dws/builder/base_builder.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Generic, Literal, TypeVar, Union
 
-from nutrient_dws.builder.staged_builders import TypedWorkflowResult
+from nutrient_dws.builder.staged_builders import TypedWorkflowResult, BufferOutput, ContentOutput, JsonContentOutput
 from nutrient_dws.http import (
     AnalyzeBuildRequestData,
     BuildRequestData,
@@ -24,7 +24,6 @@ async def _send_request(
         self,
         path: Union[Literal["/build"], Literal["/analyze_build"]],
         options: Union[BuildRequestData, AnalyzeBuildRequestData],
-        response_type: str = "json",
     ) -> Any:
         """Sends a request to the API."""
         config: RequestConfig[Union[BuildRequestData, AnalyzeBuildRequestData]] = {
@@ -34,7 +33,7 @@ async def _send_request(
             "headers": None
         }
 
-        response = await send_request(config, self.client_options, response_type)
+        response: Any = await send_request(config, self.client_options)
         return response["data"]
 
     @abstractmethod

src/nutrient_dws/builder/builder.py

@@ -500,18 +500,10 @@ async def execute(
 
             files = await self._prepare_files()
 
-            # Determine response type
-            response_type = "arraybuffer"
-            if output_config["type"] == "json-content":
-                response_type = "json"
-            elif output_config["type"] in ["html", "markdown"]:
-                response_type = "text"
-
             # Make the request
             response = await self._send_request(
                 "/build",
                 BuildRequestData(instructions=self.build_instructions, files=files),
-                response_type,
             )
 
             # Step 3: Process response
@@ -574,8 +566,7 @@ async def dry_run(self) -> WorkflowDryRunResult:
 
             response = await self._send_request(
                 "/analyze_build",
-                AnalyzeBuildRequestData(instructions=self.build_instructions),
-                "json",
+                AnalyzeBuildRequestData(instructions=self.build_instructions)
             )
 
             result["success"] = True

src/nutrient_dws/client.py

@@ -129,18 +129,17 @@ async def get_account_info(self) -> dict[str, Any]:
             print(account_info['subscriptionType'])
             ```
         """
-        response = await send_request(
+        response: Any = await send_request(
             {
                 "method": "GET",
                 "endpoint": "/account/info",
-                "data": {},
+                "data": None,
                 "headers": None,
             },
             self.options,
-            "json",
         )
 
-        return response["data"]
+        return cast(dict[str, Any], response["data"])
 
     async def create_token(self, params: CreateAuthTokenParameters) -> CreateAuthTokenResponse:
         """Create a new authentication token.
@@ -160,17 +159,17 @@ async def create_token(self, params: CreateAuthTokenParameters) -> CreateAuthTok
             print(token['id'])
             ```
         """
-        response = await send_request(
+        response: Any = await send_request(
             {
                 "method": "POST",
                 "endpoint": "/tokens",
                 "data": params,
+                "headers": None,
             },
             self.options,
-            "json",
         )
 
-        return response["data"]
+        return cast(CreateAuthTokenResponse, response["data"])
 
     async def delete_token(self, token_id: str) -> None:
         """Delete an authentication token.
@@ -187,10 +186,10 @@ async def delete_token(self, token_id: str) -> None:
             {
                 "method": "DELETE",
                 "endpoint": "/tokens",
-                "data": {"id": token_id},
+                "data": cast(Any, {"id": token_id}),
+                "headers": None,
             },
             self.options,
-            "json",
         )
 
     def workflow(self, override_timeout: Optional[int] = None) -> WorkflowInitialStage:
@@ -218,7 +217,7 @@ def workflow(self, override_timeout: Optional[int] = None) -> WorkflowInitialSta
         return StagedWorkflowBuilder(options)
 
     def _process_typed_workflow_result(
-        self, result: TypedWorkflowResult[Union[BufferOutput, ContentOutput, JsonContentOutput]]
+        self, result: TypedWorkflowResult
     ) -> Union[BufferOutput, ContentOutput, JsonContentOutput]:
         """Helper function that takes a TypedWorkflowResult, throws any errors, and returns the specific output type.
 
@@ -325,15 +324,14 @@ async def sign(
         if normalized_graphic_image:
             request_data["graphicImage"] = normalized_graphic_image
 
-        response = await send_request(
+        response: Any = await send_request(
             {
                 "method": "POST",
                 "endpoint": "/sign",
-                "data": request_data,
+                "data": cast(Any, request_data),
                 "headers": None,
             },
             self.options,
-            "arraybuffer",
         )
 
         buffer = response["data"]
@@ -372,12 +370,12 @@ async def watermark_text(
                 f.write(pdf_buffer)
             ```
         """
-        watermark_action = BuildActions.watermarkText(text, options or {})
+        watermark_action = BuildActions.watermarkText(text, cast(Any, options))
 
         builder = self.workflow().add_file_part(file, None, [watermark_action])
 
         result = await builder.output_pdf().execute()
-        return self._process_typed_workflow_result(result)
+        return cast(BufferOutput, self._process_typed_workflow_result(result))
 
     async def watermark_image(
         self,
@@ -406,12 +404,12 @@ async def watermark_image(
             pdf_buffer = result['buffer']
             ```
         """
-        watermark_action = BuildActions.watermarkImage(image, options or {})
+        watermark_action = BuildActions.watermarkImage(image, cast(Any, options))
 
         builder = self.workflow().add_file_part(file, None, [watermark_action])
 
         result = await builder.output_pdf().execute()
-        return self._process_typed_workflow_result(result)
+        return cast(BufferOutput, self._process_typed_workflow_result(result))
 
     async def convert(
         self,
@@ -496,7 +494,7 @@ async def ocr(
         builder = self.workflow().add_file_part(file, None, [ocr_action])
 
         result = await builder.output_pdf().execute()
-        return self._process_typed_workflow_result(result)
+        return cast(BufferOutput, self._process_typed_workflow_result(result))
 
     async def extract_text(
         self,
@@ -527,7 +525,7 @@ async def extract_text(
         """
         normalized_pages = normalize_page_params(pages) if pages else None
 
-        part_options = {"pages": normalized_pages} if normalized_pages else None
+        part_options = cast(Any, {"pages": normalized_pages}) if normalized_pages else None
 
         result = (
             await self.workflow()
@@ -536,7 +534,7 @@ async def extract_text(
             .execute()
         )
 
-        return self._process_typed_workflow_result(result)
+        return cast(JsonContentOutput, self._process_typed_workflow_result(result))
 
     async def extract_table(
         self,
@@ -568,7 +566,7 @@ async def extract_table(
         """
         normalized_pages = normalize_page_params(pages) if pages else None
 
-        part_options = {"pages": normalized_pages} if normalized_pages else None
+        part_options = cast(Any, {"pages": normalized_pages}) if normalized_pages else None
 
         result = (
             await self.workflow()
@@ -577,7 +575,7 @@ async def extract_table(
             .execute()
         )
 
-        return self._process_typed_workflow_result(result)
+        return cast(JsonContentOutput, self._process_typed_workflow_result(result))
 
     async def extract_key_value_pairs(
         self,
@@ -609,7 +607,7 @@ async def extract_key_value_pairs(
         """
         normalized_pages = normalize_page_params(pages) if pages else None
 
-        part_options = {"pages": normalized_pages} if normalized_pages else None
+        part_options = cast(Any, {"pages": normalized_pages}) if normalized_pages else None
 
         result = (
             await self.workflow()
@@ -618,7 +616,7 @@ async def extract_key_value_pairs(
             .execute()
         )
 
-        return self._process_typed_workflow_result(result)
+        return cast(JsonContentOutput, self._process_typed_workflow_result(result))
 
     async def password_protect(
         self,
@@ -660,9 +658,9 @@ async def password_protect(
         if permissions:
             pdf_options["userPermissions"] = permissions
 
-        result = await self.workflow().add_file_part(file).output_pdf(pdf_options).execute()
+        result = await self.workflow().add_file_part(file).output_pdf(cast(Any, pdf_options)).execute()
 
-        return self._process_typed_workflow_result(result)
+        return cast(BufferOutput, self._process_typed_workflow_result(result))
 
     async def set_metadata(
         self,
@@ -697,10 +695,10 @@ async def set_metadata(
             raise ValidationError("Invalid pdf file", {"input": pdf})
 
         result = (
-            await self.workflow().add_file_part(pdf).output_pdf({"metadata": metadata}).execute()
+            await self.workflow().add_file_part(pdf).output_pdf(cast(Any, {"metadata": metadata})).execute()
         )
 
-        return self._process_typed_workflow_result(result)
+        return cast(BufferOutput, self._process_typed_workflow_result(result))
 
     async def merge(self, files: List[FileInput]) -> BufferOutput:
         """Merge multiple documents into a single document.
@@ -733,7 +731,7 @@ async def merge(self, files: List[FileInput]) -> BufferOutput:
             workflow_builder = workflow_builder.add_file_part(file)
 
         result = await workflow_builder.output_pdf().execute()
-        return self._process_typed_workflow_result(result)
+        return cast(BufferOutput, self._process_typed_workflow_result(result))
 
     async def flatten(
         self,
@@ -774,7 +772,7 @@ async def flatten(
             await self.workflow().add_file_part(pdf, None, [flatten_action]).output_pdf().execute()
         )
 
-        return self._process_typed_workflow_result(result)
+        return cast(BufferOutput, self._process_typed_workflow_result(result))
 
     async def create_redactions_ai(
         self,
@@ -844,17 +842,16 @@ async def create_redactions_ai(
         }
 
         if options:
-            request_data["data"]["options"] = options
+            request_data["data"]["options"] = options  # type: ignore
 
-        response = await send_request(
+        response: Any = await send_request(
             {
                 "method": "POST",
                 "endpoint": "/ai/redact",
-                "data": request_data,
+                "data": cast(Any, request_data),
                 "headers": None,
             },
             self.options,
-            "arraybuffer",
         )
 
         buffer = response["data"]