You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

29 lines
1.2 KiB

  1. from django.db.models import Count, Q
  2. from django_filters.rest_framework import FilterSet, BooleanFilter
  3. from .models import Document
  4. class DocumentFilter(FilterSet):
  5. seq_annotations__isnull = BooleanFilter(field_name='seq_annotations', method='filter_annotations')
  6. doc_annotations__isnull = BooleanFilter(field_name='doc_annotations', method='filter_annotations')
  7. seq2seq_annotations__isnull = BooleanFilter(field_name='seq2seq_annotations', method='filter_annotations')
  8. def filter_annotations(self, queryset, field_name, value):
  9. queryset = queryset.annotate(num_annotations=
  10. Count(field_name, filter=
  11. Q(**{ f"{field_name}__user": self.request.user})))
  12. should_have_annotations = not value
  13. if should_have_annotations:
  14. queryset = queryset.filter(num_annotations__gte=1)
  15. else:
  16. queryset = queryset.filter(num_annotations__lte=0)
  17. return queryset
  18. class Meta:
  19. model = Document
  20. fields = ('project', 'text', 'meta', 'created_at', 'updated_at',
  21. 'doc_annotations__label__id', 'seq_annotations__label__id',
  22. 'doc_annotations__isnull', 'seq_annotations__isnull', 'seq2seq_annotations__isnull')