diff --git a/src/documents/management/commands/document_correspondents.py b/src/documents/management/commands/document_correspondents.py index 243eb9b19..0709c49d2 100644 --- a/src/documents/management/commands/document_correspondents.py +++ b/src/documents/management/commands/document_correspondents.py @@ -1,11 +1,14 @@ +import sys + from django.core.management.base import BaseCommand -from documents.models import Document, Correspondent +from documents.models import Correspondent, Document from ...mixins import Renderable class Command(Renderable, BaseCommand): + help = """ Using the current set of correspondent rules, apply said rules to all documents in the database, effectively allowing you to back-tag all @@ -13,33 +16,67 @@ class Command(Renderable, BaseCommand): after their initial import. """.replace(" ", "") + TOO_MANY_CONTINUE = ( + "Detected {} potential correspondents for {}, so we've opted for {}") + TOO_MANY_SKIP = ( + "Detected {} potential correspondents for {}, so we're skipping it") + CHANGE_MESSAGE = ( + 'Document {}: "{}" was given the correspondent id {}: "{}"') + def __init__(self, *args, **kwargs): self.verbosity = 0 BaseCommand.__init__(self, *args, **kwargs) + def add_arguments(self, parser): + parser.add_argument( + "--use-first", + default=False, + action="store_true", + help="By default this command won't try to assign a correspondent " + "if more than one matches the document. Use this flag if " + "you'd rather it just pick the first one it finds." + ) + def handle(self, *args, **options): self.verbosity = options["verbosity"] - for document in Document.objects.all(): - # No matching correspondents, so no need to continue - if document.correspondent: - continue + for document in Document.objects.filter(correspondent__isnull=True): potential_correspondents = list( Correspondent.match_all(document.content)) + if not potential_correspondents: continue potential_count = len(potential_correspondents) + correspondent = potential_correspondents[0] - selected = potential_correspondents[0] if potential_count > 1: - message = "Detected {} potential correspondents for {}, " \ - "so we've opted for {}" - print(message.format(potential_count, document, selected)) + if not options["use_first"]: + print( + self.TOO_MANY_SKIP.format(potential_count, document), + file=sys.stderr + ) + continue + print( + self.TOO_MANY_CONTINUE.format( + potential_count, + document, + correspondent + ), + file=sys.stderr + ) - print('Tagging {} with correspondent "{}"'.format(document, - selected)) - document.correspondent = selected + document.correspondent = correspondent document.save(update_fields=("correspondent",)) + + print( + self.CHANGE_MESSAGE.format( + document.pk, + document.title, + correspondent.pk, + correspondent.name + ), + file=sys.stderr + )