Merge branch 'rhaamo-add_retagger_correspondents'

src/documents/management/commands/document_correspondents.py

@@ -0,0 +1,82 @@
+import sys
+
+from django.core.management.base import BaseCommand
+
+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
+        previously indexed documents with correspondent created (or modified)
+        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.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]
+
+            if potential_count > 1:
+                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
+                )
+
+            document.correspondent = correspondent
+            document.save(update_fields=("correspondent",))
+
+            print(
+                self.CHANGE_MESSAGE.format(
+                    document.pk,
+                    document.title,
+                    correspondent.pk,
+                    correspondent.name
+                ),
+                file=sys.stderr
+            )