Allow for assignment to submissions filtered by having an attachment

diff --git a/wtem/management/commands/wtem_assign_submissions.py b/wtem/management/commands/wtem_assign_submissions.py
index 2ec816e..6d33ec8 100644 (file)
--- a/wtem/management/commands/wtem_assign_submissions.py
+++ b/wtem/management/commands/wtem_assign_submissions.py
@@ -1,24 +1,39 @@
+from optparse import make_option
+

 from django.core.management.base import BaseCommand, CommandError
 from django.db.models import Count
 from django.contrib.auth.models import User
 
 from contact.models import Contact
 from django.core.management.base import BaseCommand, CommandError
 from django.db.models import Count
 from django.contrib.auth.models import User
 
 from contact.models import Contact

-from wtem.models import Submission
+from wtem.models import Submission, Attachment

 
 
 class Command(BaseCommand):
 
 
 
 class Command(BaseCommand):
 

+    option_list = BaseCommand.option_list + (
+        make_option('--with-attachments-only',
+            action='store_true',
+            dest='attachments_only',
+            default=False,
+            help='Take into account only submissions with attachments'),
+        )
+

     def handle(self, *args, **options):
         how_many = int(args[0])
         examiner_names = args[1:]
 
     def handle(self, *args, **options):
         how_many = int(args[0])
         examiner_names = args[1:]
 

+

         users = User.objects.filter(username__in = examiner_names)
         submissions_query = Submission.objects.annotate(examiners_count = Count('examiners'))
 
         submissions = submissions_query \
         users = User.objects.filter(username__in = examiner_names)
         submissions_query = Submission.objects.annotate(examiners_count = Count('examiners'))
 
         submissions = submissions_query \

-            .filter(examiners_count__lt=2).exclude(answers = None)[0:how_many]
+            .filter(examiners_count__lt=2).exclude(answers = None)
+        
+        if options['attachments_only']:
+            with_attachment_ids = Attachment.objects.values_list('submission_id', flat=True).all()
+            submissions = submissions.filter(id__in = with_attachment_ids)

 
 

-        for submission in submissions:
+        for submission in submissions[0:how_many]:

             submission.examiners.add(*users)
             submission.save()
             self.stdout.write('added to %s:%s' % (submission.id, submission.email))
             submission.examiners.add(*users)
             submission.save()
             self.stdout.write('added to %s:%s' % (submission.id, submission.email))