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.

31 lines
1.4 KiB

  1. from django.db.models import Count, Q
  2. from django_filters.rest_framework import BooleanFilter, FilterSet
  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. speech2text_annotations__isnull = BooleanFilter(field_name='speech2text_annotations', method='filter_annotations')
  9. def filter_annotations(self, queryset, field_name, value):
  10. queryset = queryset.annotate(num_annotations=
  11. Count(field_name, filter=
  12. Q(**{ f"{field_name}__user": self.request.user}) | Q(project__collaborative_annotation=True)))
  13. should_have_annotations = not value
  14. if should_have_annotations:
  15. queryset = queryset.filter(num_annotations__gte=1)
  16. else:
  17. queryset = queryset.filter(num_annotations__lte=0)
  18. return queryset
  19. class Meta:
  20. model = Document
  21. fields = ('project', 'text', 'created_at', 'updated_at',
  22. 'doc_annotations__label__id', 'seq_annotations__label__id',
  23. 'doc_annotations__isnull', 'seq_annotations__isnull', 'seq2seq_annotations__isnull',
  24. 'speech2text_annotations__isnull')