feat: Add type hinting throughout the project

Author: qcoumes

.github/workflows/tests_and_publish.yml

@@ -84,6 +84,23 @@ jobs:
         run: |
           pydocstyle --count django_opensearch_dsl/ tests/
 
+  mypy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Setup Python
+        uses: actions/setup-python@master
+        with:
+          python-version: '3.13'
+
+      - name: Install packages
+        run: pip install -r requirements_dev.txt
+
+      - name: Mypy
+        run: |
+          mypy django_opensearch_dsl --disallow-untyped-def
+
   bandit:
     runs-on: ubuntu-latest
     steps:
@@ -130,7 +147,6 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
-
       - name: Run Opensearch in docker
         run: |
           docker compose up -d opensearch_test_${{ matrix.opensearch-version }}

.gitignore

@@ -4,6 +4,9 @@ __pycache__
 coverage.xml
 venv/
 
+# Mac
+.DS_Store
+
 # C extensions
 *.so
 

bin/pre_commit.sh

@@ -55,6 +55,21 @@ fi
 echo ""
 
 
+################################################################################
+#                                   MYPY                                       #
+################################################################################
+echo -n "${Cyan}Running mypy... $Color_Off"
+out=$(mypy django_opensearch_dsl --disallow-untyped-def)
+if [ "$?" -ne 0 ] ; then
+  echo "${Red}Error !$Color_Off"
+  echo -e "$out"
+  EXIT_CODE=1
+else
+  echo "${Green}Ok ✅ $Color_Off"
+fi
+echo ""
+
+
 ################################################################################
 #                                  BANDIT                                      #
 ################################################################################

django_opensearch_dsl/__init__.py

@@ -6,6 +6,6 @@
 __version__ = "0.7.0"
 
 
-def autodiscover():
+def autodiscover() -> None:
     """Force the import of the `documents` modules of each `INSTALLED_APPS`."""
     autodiscover_modules("documents")

django_opensearch_dsl/apps.py

@@ -1,8 +1,14 @@
+import json
+from typing import TYPE_CHECKING, Any
+
 from django.apps import AppConfig
 from django.conf import settings
 from django.utils.module_loading import import_string
 from opensearchpy.connection.connections import connections
 
+if TYPE_CHECKING:
+    from .signals import BaseSignalProcessor
+
 
 class DODConfig(AppConfig):
     """Django Opensearch DSL Appconfig."""
@@ -11,7 +17,7 @@ class DODConfig(AppConfig):
     verbose_name = "django-opensearch-dsl"
     signal_processor = None
 
-    def ready(self):
+    def ready(self) -> None:
         """Autodiscover documents and register signals."""
         self.module.autodiscover()
         connections.configure(**settings.OPENSEARCH_DSL)
@@ -21,40 +27,40 @@ def ready(self):
             self.signal_processor = self.signal_processor_class()(connections)
 
     @classmethod
-    def autosync_enabled(cls):
+    def autosync_enabled(cls) -> bool:
         """Return whether auto sync is enabled."""
         return getattr(settings, "OPENSEARCH_DSL_AUTOSYNC", True)
 
     @classmethod
-    def default_index_settings(cls):
+    def default_index_settings(cls) -> dict[str, Any]:
         """Return `OPENSEARCH_DSL_INDEX_SETTINGS`."""
         return getattr(settings, "OPENSEARCH_DSL_INDEX_SETTINGS", {})
 
     @classmethod
-    def auto_refresh_enabled(cls):
+    def auto_refresh_enabled(cls) -> bool:
         """Return whether auto refresh is enabled."""
         return getattr(settings, "OPENSEARCH_DSL_AUTO_REFRESH", False)
 
     @classmethod
-    def parallel_enabled(cls):
+    def parallel_enabled(cls) -> bool:
         """Return whether parallel operation is enabled."""
         return getattr(settings, "OPENSEARCH_DSL_PARALLEL", False)
 
     @classmethod
-    def default_queryset_pagination(cls):
+    def default_queryset_pagination(cls) -> int:
         """Return `OPENSEARCH_DSL_QUERYSET_PAGINATION`."""
         return getattr(settings, "OPENSEARCH_DSL_QUERYSET_PAGINATION", 4096)
 
     @classmethod
-    def signal_processor_class(cls):
+    def signal_processor_class(cls) -> type["BaseSignalProcessor"]:
         """Import and return the target of `OPENSEARCH_SIGNAL_PROCESSOR_CLASS`."""
         path = getattr(
             settings, "OPENSEARCH_DSL_SIGNAL_PROCESSOR", "django_opensearch_dsl.signals.RealTimeSignalProcessor"
         )
         return import_string(path)
 
     @classmethod
-    def signal_processor_serializer_class(cls):
+    def signal_processor_serializer_class(cls) -> type[json.JSONEncoder]:
         """Import and return the target of `OPENSEARCH_DSL_SIGNAL_PROCESSOR_SERIALIZER_CLASS`."""
         path = getattr(
             settings,
@@ -64,7 +70,7 @@ def signal_processor_serializer_class(cls):
         return import_string(path)
 
     @classmethod
-    def signal_processor_deserializer_class(cls):
+    def signal_processor_deserializer_class(cls) -> type[json.JSONEncoder]:
         """Import and return the target of `OPENSEARCH_DSL_SIGNAL_PROCESSOR_SERIALIZER_CLASS`."""
         path = getattr(
             settings,

django_opensearch_dsl/documents.py

@@ -1,23 +1,23 @@
-import io
 import sys
 import time
 from collections import deque
 from functools import partial
-from typing import Iterable, Optional
+from typing import Any, Callable, Iterable, Optional, TextIO, Union
 
+import opensearchpy
 from django.db import models
 from django.db.models import Q, QuerySet
 from opensearchpy.helpers import bulk, parallel_bulk
 from opensearchpy.helpers.document import Document as DSLDocument
 
 from . import fields
 from .apps import DODConfig
+from .enums import BulkAction, CommandAction
 from .exceptions import ModelFieldNotMappedError
-from .management.enums import OpensearchAction
 from .search import Search
 from .signals import post_index
 
-model_field_class_to_field_class = {
+model_field_class_to_field_class: dict[type[models.Field], type[fields.DODField]] = {
     models.AutoField: fields.IntegerField,
     models.BigAutoField: fields.LongField,
     models.BigIntegerField: fields.LongField,
@@ -48,31 +48,26 @@
 class Document(DSLDocument):
     """Allow the definition of Opensearch' index using Django `Model`."""
 
-    _prepared_fields = []
+    _prepared_fields: list[tuple[str, fields.DODField, Callable[[models.Model], Any]]] = []
 
-    def __init__(self, related_instance_to_ignore=None, **kwargs):
+    def __init__(self, related_instance_to_ignore: Any = None, **kwargs: Any) -> None:
         super(Document, self).__init__(**kwargs)
         # related instances to ignore is required to remove the instance
         # from related models on deletion.
         self._related_instance_to_ignore = related_instance_to_ignore
         self._prepared_fields = self.init_prepare()
 
     @classmethod
-    def search(cls, using=None, index=None):
-        """Return a `Search` object parametrized with the index' information."""
+    def search(cls, using: str = None, index: str = None) -> opensearchpy.Search:
+        """Return a `Search` object parametrized with the index information."""
         return Search(
             using=cls._get_using(using),
             index=cls._default_index(index),
             doc_type=[cls],
             model=cls.django.model,
         )
 
-    def get_queryset(
-        self,
-        filter_: Optional[Q] = None,
-        exclude: Optional[Q] = None,
-        count: int = None,
-    ) -> QuerySet:
+    def get_queryset(self, filter_: Optional[Q] = None, exclude: Optional[Q] = None, count: int = None) -> QuerySet:
         """Return the queryset that should be indexed by this doc type."""
         qs = self.django.model.objects.all()
 
@@ -85,7 +80,7 @@ def get_queryset(
 
         return qs
 
-    def _eta(self, start, done, total):  # pragma: no cover
+    def _eta(self, start: float, done: int, total: int) -> str:  # pragma: no cover
         if done == 0:
             return "~"
         eta = round((time.time() - start) / done * (total - done))
@@ -101,43 +96,44 @@ def get_indexing_queryset(
         filter_: Optional[Q] = None,
         exclude: Optional[Q] = None,
         count: int = None,
-        action: OpensearchAction = OpensearchAction.INDEX,
-        stdout: io.FileIO = sys.stdout,
+        action: CommandAction = CommandAction.INDEX,
+        stdout: TextIO = sys.stdout,
     ) -> Iterable:
         """Divide the queryset into chunks."""
         chunk_size = self.django.queryset_pagination
         qs = self.get_queryset(filter_=filter_, exclude=exclude, count=count)
         qs = qs.order_by("pk") if not qs.query.is_sliced else qs
-        count = qs.count()
+        total = qs.count()
         model = self.django.model.__name__
         action = action.present_participle.title()
 
         i = 0
         done = 0
         start = time.time()
         if verbose:
-            stdout.write(f"{action} {model}: 0% ({self._eta(start, done, count)})\r")
-        while done < count:
+            stdout.write(f"{action} {model}: 0% ({self._eta(start, done, total)})\r")
+        while done < total:
             if verbose:
-                stdout.write(f"{action} {model}: {round(i / count * 100)}% ({self._eta(start, done, count)})\r")
+                stdout.write(f"{action} {model}: {round(i / total * 100)}% ({self._eta(start, done, total)})\r")
 
             for obj in qs[i : i + chunk_size]:
                 done += 1
                 yield obj
 
-            i = min(i + chunk_size, count)
+            i = min(i + chunk_size, total)
 
         if verbose:
-            stdout.write(f"{action} {count} {model}: OK          \n")
+            stdout.write(f"{action} {total} {model}: OK          \n")
 
-    def init_prepare(self):
+    def init_prepare(self) -> list[tuple[str, fields.DODField, Callable[[models.Model], Any]]]:
         """Initialise the data model preparers once here.
 
         Extracts the preparers from the model and generate a list of callables
         to avoid doing that work on every object instance over.
         """
-        index_fields = getattr(self, "_fields", {})
+        index_fields: dict[str, fields.DODField] = getattr(self, "_fields", {})
         preparers = []
+        fn: Callable[[models.Model], Any]
         for name, field in iter(index_fields.items()):
             if not isinstance(field, fields.DODField):  # pragma: no cover
                 continue
@@ -162,13 +158,13 @@ def init_prepare(self):
 
         return preparers
 
-    def prepare(self, instance):
+    def prepare(self, instance: models.Model) -> dict[str, Any]:
         """Generate the opensearch's document from `instance` based on defined fields."""
         data = {name: prep_func(instance) for name, field, prep_func in self._prepared_fields}
         return data
 
     @classmethod
-    def to_field(cls, field_name, model_field):
+    def to_field(cls, field_name: str, model_field: models.Field) -> fields.DODField:
         """Return the opensearch field instance mapped to the model field class.
 
         This is a good place to hook into if you have more complex
@@ -179,14 +175,16 @@ def to_field(cls, field_name, model_field):
         except KeyError:  # pragma: no cover
             raise ModelFieldNotMappedError(f"Cannot convert model field {field_name} to an Opensearch field!")
 
-    def bulk(self, actions, using=None, **kwargs):
+    def bulk(
+        self, actions: Iterable[dict[str, Any]], using: str = None, **kwargs: Any
+    ) -> Union[tuple[int, int], tuple[int, list]]:
         """Execute given actions in bulk."""
         response = bulk(client=self._get_connection(using), actions=actions, **kwargs)
         # send post index signal
         post_index.send(sender=self.__class__, instance=self, actions=actions, response=response)
         return response
 
-    def parallel_bulk(self, actions, using=None, **kwargs):
+    def parallel_bulk(self, actions: Iterable[dict[str, Any]], using: str = None, **kwargs: Any) -> tuple[int, list]:
         """Parallel version of `bulk`."""
         kwargs.setdefault("chunk_size", self.django.queryset_pagination)
         bulk_actions = parallel_bulk(client=self._get_connection(using), actions=actions, **kwargs)
@@ -199,7 +197,7 @@ def parallel_bulk(self, actions, using=None, **kwargs):
         return 1, []
 
     @classmethod
-    def generate_id(cls, object_instance):
+    def generate_id(cls, object_instance: models.Model) -> Any:
         """Generate the opensearch's _id from a Django `Model` instance.
 
         The default behavior is to use the Django object's pk (id) as the
@@ -208,47 +206,52 @@ def generate_id(cls, object_instance):
         """
         return object_instance.pk
 
-    def _prepare_action(self, object_instance, action):
+    def _prepare_action(self, object_instance: models.Model, action: BulkAction) -> dict[str, Any]:
         return {
             "_op_type": action,
             "_index": self._index._name,  # noqa
             "_id": self.generate_id(object_instance),
             "_source" if action != "update" else "doc": (self.prepare(object_instance) if action != "delete" else None),
         }
 
-    def _get_actions(self, object_list, action):
+    def _get_actions(self, object_list: Iterable[models.Model], action: BulkAction) -> Iterable[dict[str, Any]]:
         for object_instance in object_list:
             if action == "delete" or self.should_index_object(object_instance):
                 yield self._prepare_action(object_instance, action)
 
-    def _bulk(self, *args, parallel=False, using=None, **kwargs):
+    def _bulk(
+        self, actions: Iterable[dict[str, Any]], parallel: bool = False, using: str = None, **kwargs: Any
+    ) -> Union[tuple[int, int], tuple[int, list]]:
         """Allow switching between normal and parallel bulk operation."""
         if parallel:
-            return self.parallel_bulk(*args, using=using, **kwargs)
-        return self.bulk(*args, using=using, **kwargs)
+            return self.parallel_bulk(actions, using=using, **kwargs)
+        return self.bulk(actions, using=using, **kwargs)
 
-    def should_index_object(self, obj):
+    def should_index_object(self, object_instance: models.Model) -> bool:
         """Whether given object should be indexed.
 
         Overwriting this method and returning a boolean value should determine
         whether the object should be indexed.
         """
         return True
 
-    def update(self, thing, action, *args, refresh=None, using=None, **kwargs):  # noqa
+    def update(  # type: ignore[override]  # noqa
+        self,
+        thing: Union[models.Model, Iterable[models.Model]],
+        action: BulkAction,
+        refresh: bool = None,
+        parallel: bool = None,
+        using: str = None,
+        **kwargs: Any,
+    ) -> Union[tuple[int, int], tuple[int, list]]:
         """Update document in OS for a model, iterable of models or queryset."""
         if refresh is None:
             refresh = getattr(self.Index, "auto_refresh", DODConfig.auto_refresh_enabled())
+        if parallel is None:
+            parallel = DODConfig.parallel_enabled()
 
-        if isinstance(thing, models.Model):
-            object_list = [thing]
-        else:
-            object_list = thing
+        object_list = [thing] if isinstance(thing, models.Model) else thing
 
         return self._bulk(
-            self._get_actions(object_list, action),
-            *args,
-            refresh=refresh,
-            using=using,
-            **kwargs,
+            self._get_actions(object_list, action), parallel=parallel, refresh=refresh, using=using, **kwargs
         )