feat: Implement get_documents() for DocumentRegistry

Signed-off-by: gcarq <[email protected]>

Author: gcarq

django_opensearch_dsl/registries.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
 from copy import deepcopy
+from itertools import chain
 
 from django.core.exceptions import ImproperlyConfigured, ObjectDoesNotExist
 from django.db.models import Model
@@ -161,6 +162,14 @@ def delete(self, instance, **kwargs):
         """
         self.update(instance, action="delete", **kwargs)
 
+    def get_documents(self, models=None):
+        """Get all documents in the registry or the documents for a list of models."""
+        if models is not None:
+            docs_iter = (self._models[model] for model in models if model in self._models)
+        else:
+            docs_iter = self._models.values()
+        return set(chain.from_iterable(docs_iter))
+
     def get_models(self):
         """Get all models in the registry."""
         return set(self._models.keys())

tests/tests/test_registries.py

@@ -57,6 +57,14 @@ def test_get_indices_by_unregister_model(self):
         ModelC = Mock()
         self.assertFalse(self.registry.get_indices([ModelC]))
 
+    def test_get_documents(self):
+        docs = {self.doc_a1, self.doc_a2, self.doc_b1, self.doc_c1, self.doc_d1, self.doc_e1}
+        self.assertEqual(self.registry.get_documents(), docs)
+
+    def test_get_documents_by_models(self):
+        docs = {self.doc_a1, self.doc_a2, self.doc_b1}
+        self.assertEqual(self.registry.get_documents(models=[self.ModelA, self.ModelB]), docs)
+
     def test_get_models(self):
         self.assertEqual(self.registry.get_models(), {self.ModelA, self.ModelB, self.ModelC, self.ModelD, self.ModelE})