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.

33 lines
1.5 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=Count(
  11. expression=field_name,
  12. filter=Q(**{f"{field_name}__user": self.request.user}) | Q(project__collaborative_annotation=True)
  13. )
  14. )
  15. should_have_annotations = not value
  16. if should_have_annotations:
  17. queryset = queryset.filter(num_annotations__gte=1)
  18. else:
  19. queryset = queryset.filter(num_annotations__lte=0)
  20. return queryset
  21. class Meta:
  22. model = Document
  23. fields = ('project', 'text', 'created_at', 'updated_at',
  24. 'doc_annotations__label__id', 'seq_annotations__label__id',
  25. 'doc_annotations__isnull', 'seq_annotations__isnull', 'seq2seq_annotations__isnull',
  26. 'speech2text_annotations__isnull')