signal_documentation/src/signals/filters.py at 7bddeb593d289227aca5849a2e5288be559aa3a8 · cmu-delphi/signal_documentation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
from typing import Any

import django_filters
from django.db.models import Q
from django_filters.filters import (
    BaseInFilter,
    CharFilter,
    NumberFilter,
    OrderingFilter,
)
from django_filters.widgets import QueryArrayWidget

from datasources.models import SourceSubdivision
from signals.models import (
    FormatChoices,
    Signal,
    TimeLabelChoices,
)


class NumberInFilter(BaseInFilter, NumberFilter):
    pass


class SignalFilter(django_filters.FilterSet):
    """
    FilterSet for the Signal model.
    """

    id = NumberInFilter(
        field_name='id',
        lookup_expr='in',
        widget=QueryArrayWidget
    )
    search = CharFilter(method='filter_search')
    order_by = OrderingFilter(
        fields=(
            ('display_name', 'name'),
            ('source__name', 'source'),
            ('last_updated', 'last_updated'),
        )
    )
    format_type = django_filters.MultipleChoiceFilter(choices=FormatChoices.choices)
    source = django_filters.ModelMultipleChoiceFilter(queryset=SourceSubdivision.objects.all())
    time_label = django_filters.MultipleChoiceFilter(choices=TimeLabelChoices.choices)

    class Meta:
        model = Signal
        fields: list[str] = [
            'id',
            'search',
            'pathogen',
            'active',
            'available_geography',
            'signal_type',
            'category',
            'format_type',
            'source',
            'time_label',
        ]

    def filter_search(self, queryset, name, value) -> Any:
        """
        Custom filter method to perform a search on the Signal model.

        Args:
            queryset (QuerySet): The initial queryset.
            name (str): The name of the filter field.
            value (Any): The value to search for.

        Returns:
            QuerySet: The filtered queryset based on the search value.
        """

        if not value:
            return queryset
        search_tokens = value.split()
        queries: list[Q] = []
        for field in ['name', 'description', 'short_description']:
            token_query: list[Q] = []
            for token in search_tokens:
                if '*' in token:
                    left = token.find('*') == 0
                    right = token.rfind('*') == len(token) - 1
                    token = token.replace('*', '')
                    if left and right:
                        token_query.append(Q((f'{field}__icontains', token)))
                        continue
                    if left:
                        token_query.append(Q((f'{field}__iregex', fr'{token}\b')))
                        continue
                    if right:
                        token_query.append(Q((f'{field}__iregex', fr'\b{token}')))
                        continue
                else:
                    token_query.append(Q((f'{field}__iregex', fr"\b{token}\b")))
                    continue
            query: Q = token_query.pop()
            for item in token_query:
                query &= item
            queries.append(query)

        query = queries.pop()

        for item in queries:
            query |= item

        return queryset.filter(query)