Django REST filter around `django.contrib.postgres.search` expressions

drf_full_text_search_filter.py

from typing import Union, Collection

from django.contrib.postgres.search import SearchVector, SearchQuery
from django.db.models.expressions import BaseExpression

from django_filters import rest_framework as filters
from django_filters.constants import EMPTY_VALUES


class FullTextSearchFilter(filters.CharFilter):
    """ Django REST filter around ``django.contrib.postgres.search`` expressions.

    Accepts ``SearchVector`` instance, string or collection of expressions,
    ``search_type`` for ``SearchQuery`` and locale ``config``

    Example:
    >>> search = FullTextSearchFilter(
    >>>     vector=("field_1", "field__subfield"),  # or `SearchVector` instance
    >>>     search_type="phrase",
    >>>     config="french",
    >>> )
    """

    _InnerExpr = Union[str, BaseExpression]
    _VectorExpr = Union[_InnerExpr, Collection[_InnerExpr], SearchVector]

    def __init__(
        self,
        *args,
        vector: _VectorExpr,
        suffix: str = "search",
        search_type: str = "plain",
        config: str = None,
        **kwargs,
    ):
        """
        :param vector: Either ``SearchVector``, string or collection of expressions.
        :param suffix: Appendix for query annotation, optional.
        :param search_type: ``SearchQuery``'s search type.
        :param config: Locale config for ``SearchVector``
        """
        if callable(vector):
            vector = vector()
        if isinstance(vector, (str, BaseExpression)):
            vector = SearchVector(vector, config=config)
        elif not isinstance(vector, SearchVector):
            vector = SearchVector(*vector, config=config)

        self.vector = vector
        self.suffix = suffix
        self.search_type = search_type
        super().__init__(*args, **kwargs)

    def filter(self, qs, value):
        if value in EMPTY_VALUES:
            return qs

        search_field = self._search_field
        query = SearchQuery(value, search_type=self.search_type)

        annotation, lookup = {search_field: self.vector}, {search_field: query}
        return self.get_method(qs.annotate(**annotation))(**lookup)

    @property
    def _search_field(self):
        return f"{self.field_name}_{self.suffix}"