fix: resolve typesense imports in mypy and fix the type errors it surfaces

Author: tharropoulos

setup.cfg

@@ -56,6 +56,7 @@ enable_error_code =
   redundant-self,
 
 explicit_package_bases = true
+mypy_path = src
 ignore_missing_imports = true
 strict = true
 warn_unreachable = true

src/typesense/async_/analytics_rule_v1.py

@@ -74,11 +74,9 @@ async def retrieve(
             Union[RuleSchemaForQueries, RuleSchemaForCounters]:
               The schema containing the rule details.
         """
-        response: typing.Union[
-            RuleSchemaForQueries, RuleSchemaForCounters
-        ] = await self.api_call.get(
+        response = await self.api_call.get(
             self._endpoint_path,
-            entity_type=dict,
+            entity_type=typing.Dict[str, typing.Any],
             as_json=True,
         )
         return typing.cast(
@@ -101,7 +99,7 @@ async def delete(self) -> RuleDeleteSchema:
         return response
 
     @property
-    @warn_deprecation(  # type: ignore[untyped-decorator]
+    @warn_deprecation(
         "AsyncAnalyticsRuleV1 is deprecated on v30+. Use client.analytics.rules[rule_id] instead.",
         flag_name="analytics_rules_v1_deprecation",
     )

src/typesense/async_/analytics_rules_v1.py

@@ -89,7 +89,7 @@ def __getitem__(self, rule_id: str) -> AsyncAnalyticsRuleV1:
             self.rules[rule_id] = AsyncAnalyticsRuleV1(self.api_call, rule_id)
         return self.rules[rule_id]
 
-    @warn_deprecation(  # type: ignore[untyped-decorator]
+    @warn_deprecation(
         "AsyncAnalyticsRulesV1 is deprecated on v30+. Use client.analytics instead.",
         flag_name="analytics_rules_v1_deprecation",
     )
@@ -115,21 +115,19 @@ async def create(
                 The created rule. Returns RuleSchemaForCounters for counter rules
                 and RuleSchemaForQueries for query rules.
         """
-        response: typing.Union[
-            RuleSchemaForCounters, RuleSchemaForQueries
-        ] = await self.api_call.post(
+        response = await self.api_call.post(
             AsyncAnalyticsRulesV1.resource_path,
             body=rule,
             params=rule_parameters,
             as_json=True,
-            entity_type=dict,
+            entity_type=typing.Dict[str, typing.Any],
         )
         return typing.cast(
             typing.Union[RuleSchemaForCounters, RuleSchemaForQueries],
             response,
         )
 
-    @warn_deprecation(  # type: ignore[untyped-decorator]
+    @warn_deprecation(
         "AsyncAnalyticsRulesV1 is deprecated on v30+. Use client.analytics instead.",
         flag_name="analytics_rules_v1_deprecation",
     )
@@ -148,19 +146,17 @@ async def upsert(
         Returns:
             Union[RuleSchemaForCounters, RuleCreateSchemaForQueries]: The upserted rule.
         """
-        response: typing.Union[
-            RuleSchemaForCounters, RuleCreateSchemaForQueries
-        ] = await self.api_call.put(
+        response = await self.api_call.put(
             "/".join([self.resource_path, rule_id]),
             body=rule,
-            entity_type=dict,
+            entity_type=typing.Dict[str, typing.Any],
         )
         return typing.cast(
             typing.Union[RuleSchemaForCounters, RuleCreateSchemaForQueries],
             response,
         )
 
-    @warn_deprecation(  # type: ignore[untyped-decorator]
+    @warn_deprecation(
         "AsyncAnalyticsRulesV1 is deprecated on v30+. Use client.analytics instead.",
         flag_name="analytics_rules_v1_deprecation",
     )

src/typesense/async_/api_call.py

@@ -60,7 +60,7 @@
     import typing_extensions as typing
 
 TEntityDict = typing.TypeVar("TEntityDict")
-TParams = typing.TypeVar("TParams", bound=typing.Dict[str, typing.Any])
+TParams = typing.TypeVar("TParams", bound=typing.Mapping[str, object])
 TBody = typing.TypeVar(
     "TBody", bound=typing.Union[str, bytes, typing.Mapping[str, typing.Any]]
 )
@@ -95,7 +95,7 @@ class SessionFunctionKwargs(typing.Generic[TParams, TBody], typing.TypedDict):
 
     params: typing.NotRequired[typing.Union[TParams, None]]
     data: typing.NotRequired[typing.Union[TBody, None]]
-    content: typing.NotRequired[typing.Union[TBody, str, None]]
+    content: typing.NotRequired[typing.Union[str, bytes, None]]
     headers: typing.NotRequired[typing.Dict[str, str]]
     timeout: typing.NotRequired[float]
 

src/typesense/async_/client.py

@@ -164,5 +164,5 @@ def typed_collection(
         """
         if name is None:
             name = model.__name__.lower()
-        collection: AsyncCollection[TDoc] = self.collections[name]
-        return collection
+        # ``collections`` is typed for the default DocumentSchema; narrow it to the model.
+        return typing.cast(AsyncCollection[TDoc], self.collections[name])

src/typesense/async_/documents.py

@@ -61,6 +61,17 @@
     None,
 ]
 
+# One line of an import response. ``ImportResponse`` is a union of lists, one per
+# return mode, so the helpers below build a list of these and ``import_`` casts it
+# to the list type its overloads promise.
+_ImportResponseItem = typing.Union[
+    ImportResponseWithDoc[TDoc],
+    ImportResponseWithId,
+    ImportResponseWithDocAndId[TDoc],
+    ImportResponseSuccess,
+    ImportResponseFail[TDoc],
+]
+
 
 class AsyncDocuments(typing.Generic[TDoc]):
     """
@@ -125,12 +136,14 @@ async def create(
         Returns:
             TDoc: The created document.
         """
-        dirty_values_parameters = dirty_values_parameters or {}
-        dirty_values_parameters["action"] = "create"
+        write_parameters: typing.Dict[str, object] = {
+            **(dirty_values_parameters or {}),
+            "action": "create",
+        }
         response = await self.api_call.post(
             self._endpoint_path(),
             body=document,
-            params=dirty_values_parameters,
+            params=write_parameters,
             as_json=True,
             entity_type=typing.Dict[str, str],
         )
@@ -154,7 +167,14 @@ async def create_many(
                 The list of import responses.
         """
         logger.warn("`create_many` is deprecated: please use `import_`.")
-        return await self.import_(documents, dirty_values_parameters)
+        # Dirty values parameters are a subset of the write parameters.
+        return await self.import_(
+            documents,
+            typing.cast(
+                typing.Optional[DocumentWriteParameters],
+                dirty_values_parameters,
+            ),
+        )
 
     async def upsert(
         self,
@@ -172,12 +192,14 @@ async def upsert(
         Returns:
             TDoc: The upserted document.
         """
-        dirty_values_parameters = dirty_values_parameters or {}
-        dirty_values_parameters["action"] = "upsert"
+        write_parameters: typing.Dict[str, object] = {
+            **(dirty_values_parameters or {}),
+            "action": "upsert",
+        }
         response = await self.api_call.post(
             self._endpoint_path(),
             body=document,
-            params=dirty_values_parameters,
+            params=write_parameters,
             as_json=True,
             entity_type=typing.Dict[str, str],
         )
@@ -199,12 +221,14 @@ async def update(
         Returns:
             UpdateByFilterResponse: The response containing information about the update.
         """
-        dirty_values_parameters = dirty_values_parameters or {}
-        dirty_values_parameters["action"] = "update"
+        update_parameters: typing.Dict[str, object] = {
+            **(dirty_values_parameters or {}),
+            "action": "update",
+        }
         response: UpdateByFilterResponse = await self.api_call.patch(
             self._endpoint_path(),
             body=document,
-            params=dirty_values_parameters,
+            params=update_parameters,
             entity_type=UpdateByFilterResponse,
         )
         return response
@@ -301,9 +325,14 @@ async def import_(
             return await self._import_raw(documents, import_parameters)
 
         if batch_size:
-            return await self._batch_import(documents, import_parameters, batch_size)
-
-        return await self._bulk_import(documents, import_parameters)
+            response_objs = await self._batch_import(
+                documents,
+                import_parameters,
+                batch_size,
+            )
+        else:
+            response_objs = await self._bulk_import(documents, import_parameters)
+        return typing.cast(ImportResponse[TDoc], response_objs)
 
     async def export(
         self,
@@ -410,9 +439,9 @@ async def _batch_import(
         documents: typing.List[TDoc],
         import_parameters: _ImportParameters,
         batch_size: int,
-    ) -> ImportResponse[TDoc]:
+    ) -> typing.List[_ImportResponseItem[TDoc]]:
         """Import documents in batches."""
-        response_objs: ImportResponse[TDoc] = []
+        response_objs: typing.List[_ImportResponseItem[TDoc]] = []
         for batch_index in range(0, len(documents), batch_size):
             batch = documents[batch_index : batch_index + batch_size]
             api_response = await self._bulk_import(batch, import_parameters)
@@ -423,7 +452,7 @@ async def _bulk_import(
         self,
         documents: typing.List[TDoc],
         import_parameters: _ImportParameters,
-    ) -> ImportResponse[TDoc]:
+    ) -> typing.List[_ImportResponseItem[TDoc]]:
         """Import a list of documents in bulk."""
         document_strs = [json.dumps(doc) for doc in documents]
         if not document_strs:
@@ -439,9 +468,12 @@ async def _bulk_import(
         )
         return self._parse_import_response(res)
 
-    def _parse_import_response(self, response: str) -> ImportResponse[TDoc]:
+    def _parse_import_response(
+        self,
+        response: str,
+    ) -> typing.List[_ImportResponseItem[TDoc]]:
         """Parse the import response string into a list of response objects."""
-        response_objs: typing.List[ImportResponse] = []
+        response_objs: typing.List[_ImportResponseItem[TDoc]] = []
         for res_obj_str in response.split("\n"):
             try:
                 res_obj_json = json.loads(res_obj_str)

src/typesense/async_/keys.py

@@ -103,11 +103,11 @@ async def create(self, schema: ApiKeyCreateSchema) -> ApiKeyCreateResponseSchema
             ...     }
             ... )
         """
-        response: ApiKeySchema = await self.api_call.post(
+        response: ApiKeyCreateResponseSchema = await self.api_call.post(
             AsyncKeys.resource_path,
             as_json=True,
             body=schema,
-            entity_type=ApiKeySchema,
+            entity_type=ApiKeyCreateResponseSchema,
         )
         return response
 

src/typesense/async_/operations.py

@@ -60,8 +60,10 @@ def __init__(self, api_call: AsyncApiCall):
         """
         self.api_call = api_call
 
+    # The generic ``str`` overload below also matches "schema_changes"; overloads are
+    # tried in order, so this one wins.
     @typing.overload
-    async def perform(
+    async def perform(  # type: ignore[overload-overlap]
         self,
         operation_name: typing.Literal["schema_changes"],
         query_params: None = None,
@@ -132,36 +134,36 @@ async def perform(
     @typing.overload
     async def perform(
         self,
-        operation_name: str,
-        query_params: typing.Union[typing.Dict[str, str], None] = None,
+        operation_name: typing.Literal["snapshot"],
+        query_params: SnapshotParameters,
     ) -> OperationResponse:
         """
-        Perform a generic operation.
+        Perform a snapshot operation.
 
         Args:
-            operation_name (str): The name of the operation.
-            query_params (Union[Dict[str, str], None], optional):
-                Query parameters for the operation.
+            operation_name (Literal["snapshot"]): The name of the operation.
+            query_params (SnapshotParameters): Query parameters for the snapshot operation.
 
         Returns:
-            OperationResponse: The response from the operation.
+            OperationResponse: The response from the snapshot operation.
         """
 
     @typing.overload
     async def perform(
         self,
-        operation_name: typing.Literal["snapshot"],
-        query_params: SnapshotParameters,
+        operation_name: str,
+        query_params: typing.Union[typing.Dict[str, str], None] = None,
     ) -> OperationResponse:
         """
-        Perform a snapshot operation.
+        Perform a generic operation.
 
         Args:
-            operation_name (Literal["snapshot"]): The name of the operation.
-            query_params (SnapshotParameters): Query parameters for the snapshot operation.
+            operation_name (str): The name of the operation.
+            query_params (Union[Dict[str, str], None], optional):
+                Query parameters for the operation.
 
         Returns:
-            OperationResponse: The response from the snapshot operation.
+            OperationResponse: The response from the operation.
         """
 
     async def perform(
@@ -181,7 +183,7 @@ async def perform(
             typing.Dict[str, str],
             None,
         ] = None,
-    ) -> OperationResponse:
+    ) -> typing.Union[OperationResponse, typing.List[SchemaChangesResponse]]:
         """
         Perform an operation on the Typesense API.
 
@@ -202,13 +204,16 @@ async def perform(
             >>> response = await operations.perform("vote")
             >>> health = await operations.is_healthy()
         """
-        response: OperationResponse = await self.api_call.post(
+        response = await self.api_call.post(
             self._endpoint_path(operation_name),
             params=query_params,
             as_json=True,
-            entity_type=OperationResponse,
+            entity_type=object,
+        )
+        return typing.cast(
+            typing.Union[OperationResponse, typing.List[SchemaChangesResponse]],
+            response,
         )
-        return response
 
     async def is_healthy(self) -> bool:
         """
@@ -222,16 +227,14 @@ async def is_healthy(self) -> bool:
             >>> healthy = await operations.is_healthy()
             >>> print(healthy)
         """
-        call_resp: HealthCheckResponse = await self.api_call.get(
+        call_resp: object = await self.api_call.get(
             AsyncOperations.health_path,
             as_json=True,
             entity_type=HealthCheckResponse,
         )
-        if isinstance(call_resp, typing.Dict):
-            is_ok: bool = call_resp.get("ok", False)
-        else:
-            is_ok = False
-        return is_ok
+        if isinstance(call_resp, dict):
+            return bool(call_resp.get("ok", False))
+        return False
 
     async def toggle_slow_request_log(
         self,

src/typesense/async_/override.py

@@ -87,7 +87,7 @@ async def delete(self) -> OverrideDeleteSchema:
         return response
 
     @property
-    @warn_deprecation(  # type: ignore[untyped-decorator]
+    @warn_deprecation(
         "The override API (collections/{collection}/overrides/{override_id}) is deprecated is removed on v30+. "
         "Use curation sets (curation_sets) instead.",
         flag_name="overrides_deprecation",

src/typesense/async_/overrides.py

@@ -129,7 +129,7 @@ async def retrieve(self) -> OverrideRetrieveSchema:
         )
         return response
 
-    @warn_deprecation(  # type: ignore[untyped-decorator]
+    @warn_deprecation(
         "AsyncOverrides is deprecated on v30+. Use client.curation_sets instead.",
         flag_name="overrides_deprecation",
     )