From e23db3a77b7ca3e4636609a055fffb754eadaa5a Mon Sep 17 00:00:00 2001
From: Silvano Cerza <silvanocerza@gmail.com>
Date: Mon, 2 Oct 2023 08:36:41 -0700
Subject: [PATCH] Fix mypy issues

---
 src/elasticsearch_haystack/document_store.py | 7 +++++--
 src/elasticsearch_haystack/filters.py        | 4 ++--
 tests/test_document_store.py                 | 2 +-
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/elasticsearch_haystack/document_store.py b/src/elasticsearch_haystack/document_store.py
index 9b08fcf..cf7f6c0 100644
--- a/src/elasticsearch_haystack/document_store.py
+++ b/src/elasticsearch_haystack/document_store.py
@@ -161,11 +161,14 @@ def write_documents(self, documents: List[Document], policy: DuplicatePolicy = D
             index=self._index,
             raise_on_error=False,
         )
-
         if errors and policy == DuplicatePolicy.FAIL:
             # TODO: Handle errors in a better way, we're assuming that all errors
             # are related to duplicate documents but that could be very well be wrong.
-            ids = ", ".join(e["create"]["_id"] for e in errors)
+
+            # mypy complains that `errors`` could be either `int` or a `list` of `dict`s.
+            # Since the type depends on the parameters passed to `helpers.bulk()`` we know
+            # for sure that it will be a `list`.
+            ids = ", ".join(e["create"]["_id"] for e in errors)  # type: ignore[union-attr]
             msg = f"IDs '{ids}' already exist in the document store."
             raise DuplicateDocumentError(msg)
 
diff --git a/src/elasticsearch_haystack/filters.py b/src/elasticsearch_haystack/filters.py
index 66041fa..31f111e 100644
--- a/src/elasticsearch_haystack/filters.py
+++ b/src/elasticsearch_haystack/filters.py
@@ -39,7 +39,7 @@ def _normalize_filters(filters: Union[List[Dict], Dict], logical_condition="") -
 
 
 def _parse_comparison(field: str, comparison: Union[Dict, List, str, float]) -> List:
-    result = []
+    result: List[Dict[str, Any]] = []
     if isinstance(comparison, dict):
         for comparator, val in comparison.items():
             if comparator == "$eq":
@@ -125,7 +125,7 @@ def _normalize_ranges(conditions: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     range_conditions = [next(iter(c["range"].items())) for c in conditions if "range" in c]
     if range_conditions:
         conditions = [c for c in conditions if "range" not in c]
-        range_conditions_dict = {}
+        range_conditions_dict: Dict[str, Any] = {}
         for field_name, comparison in range_conditions:
             if field_name not in range_conditions_dict:
                 range_conditions_dict[field_name] = {}
diff --git a/tests/test_document_store.py b/tests/test_document_store.py
index 8dd7fcd..35f566c 100644
--- a/tests/test_document_store.py
+++ b/tests/test_document_store.py
@@ -17,7 +17,7 @@ class TestDocumentStore(DocumentStoreBaseTests):
     """
 
     @pytest.fixture
-    def docstore(self, request) -> ElasticsearchDocumentStore:
+    def docstore(self, request):
         """
         This is the most basic requirement for the child class: provide
         an instance of this document store so the base class can use it.