refactor

src/documents/consumer.py

@@ -30,19 +30,19 @@ class Consumer:
 
         self.logger = logging.getLogger(__name__)
         self.logging_group = None
+        self.path = None
+        self.filename = None
+        self.override_title = None
+        self.override_correspondent_id = None
+        self.override_tag_ids = None
+        self.override_document_type_id = None
 
-        self.storage_type = Document.STORAGE_TYPE_UNENCRYPTED
-        if settings.PASSPHRASE:
-            self.storage_type = Document.STORAGE_TYPE_GPG
-
-    @staticmethod
-    def pre_check_file_exists(filename):
-        if not os.path.isfile(filename):
+    def pre_check_file_exists(self):
+        if not os.path.isfile(self.path):
             raise ConsumerError("Cannot consume {}: It is not a file".format(
-                filename))
+                self.path))
 
-    @staticmethod
-    def pre_check_consumption_dir():
+    def pre_check_consumption_dir(self):
         if not settings.CONSUMPTION_DIR:
             raise ConsumerError(
                 "The CONSUMPTION_DIR settings variable does not appear to be "
@@ -53,26 +53,23 @@ class Consumer:
                 "Consumption directory {} does not exist".format(
                     settings.CONSUMPTION_DIR))
 
-    @staticmethod
-    def pre_check_regex(filename):
-        if not re.match(FileInfo.REGEXES["title"], filename):
+    def pre_check_regex(self):
+        if not re.match(FileInfo.REGEXES["title"], self.filename):
             raise ConsumerError(
                 "Filename {} does not seem to be safe to "
-                "consume".format(filename))
+                "consume".format(self.filename))
 
-    @staticmethod
-    def pre_check_duplicate(filename):
-        with open(filename, "rb") as f:
+    def pre_check_duplicate(self):
+        with open(self.path, "rb") as f:
             checksum = hashlib.md5(f.read()).hexdigest()
         if Document.objects.filter(checksum=checksum).exists():
             if settings.CONSUMER_DELETE_DUPLICATES:
-                os.unlink(filename)
+                os.unlink(self.path)
             raise ConsumerError(
-                "Not consuming {}: It is a duplicate.".format(filename)
+                "Not consuming {}: It is a duplicate.".format(self.filename)
             )
 
-    @staticmethod
-    def pre_check_directories():
+    def pre_check_directories(self):
         os.makedirs(settings.SCRATCH_DIR, exist_ok=True)
         os.makedirs(settings.THUMBNAIL_DIR, exist_ok=True)
         os.makedirs(settings.ORIGINALS_DIR, exist_ok=True)
@@ -83,16 +80,23 @@ class Consumer:
         })
 
     def try_consume_file(self,
-                         filename,
-                         original_filename=None,
-                         force_title=None,
-                         force_correspondent_id=None,
-                         force_document_type_id=None,
-                         force_tag_ids=None):
+                         path,
+                         override_filename=None,
+                         override_title=None,
+                         override_correspondent_id=None,
+                         override_document_type_id=None,
+                         override_tag_ids=None):
         """
         Return the document object if it was successfully created.
         """
 
+        self.path = path
+        self.filename = override_filename or os.path.basename(path)
+        self.override_title = override_title
+        self.override_correspondent_id = override_correspondent_id
+        self.override_document_type_id = override_document_type_id
+        self.override_tag_ids = override_tag_ids
+
         # this is for grouping logging entries for this particular file
         # together.
 
@@ -100,19 +104,19 @@ class Consumer:
 
         # Make sure that preconditions for consuming the file are met.
 
-        self.pre_check_file_exists(filename)
+        self.pre_check_file_exists()
         self.pre_check_consumption_dir()
         self.pre_check_directories()
-        self.pre_check_regex(filename)
-        self.pre_check_duplicate(filename)
+        self.pre_check_regex()
+        self.pre_check_duplicate()
 
-        self.log("info", "Consuming {}".format(filename))
+        self.log("info", "Consuming {}".format(self.filename))
 
         # Determine the parser class.
 
-        parser_class = get_parser_class(original_filename or filename)
+        parser_class = get_parser_class(self.filename)
         if not parser_class:
-            raise ConsumerError("No parsers abvailable for {}".format(filename))
+            raise ConsumerError("No parsers abvailable for {}".format(self.filename))
         else:
             self.log("debug", "Parser: {}".format(parser_class.__name__))
 
@@ -120,13 +124,13 @@ class Consumer:
 
         document_consumption_started.send(
             sender=self.__class__,
-            filename=filename,
+            filename=self.path,
             logging_group=self.logging_group
         )
 
         # This doesn't parse the document yet, but gives us a parser.
 
-        document_parser = parser_class(filename, self.logging_group)
+        document_parser = parser_class(self.path, self.logging_group)
 
         # However, this already created working directories which we have to
         # clean up.
@@ -134,9 +138,9 @@ class Consumer:
         # Parse the document. This may take some time.
 
         try:
-            self.log("debug", "Generating thumbnail for {}...".format(filename))
+            self.log("debug", "Generating thumbnail for {}...".format(self.filename))
             thumbnail = document_parser.get_optimised_thumbnail()
-            self.log("debug", "Parsing {}...".format(filename))
+            self.log("debug", "Parsing {}...".format(self.filename))
             text = document_parser.get_text()
             date = document_parser.get_date()
         except ParseError as e:
@@ -165,14 +169,7 @@ class Consumer:
                 # store the document.
                 document = self._store(
                     text=text,
-                    doc=filename,
-                    thumbnail=thumbnail,
-                    date=date,
-                    original_filename=original_filename,
-                    force_title=force_title,
-                    force_correspondent_id=force_correspondent_id,
-                    force_document_type_id=force_document_type_id,
-                    force_tag_ids=force_tag_ids
+                    date=date
                 )
 
                 # If we get here, it was successful. Proceed with post-consume
@@ -189,12 +186,12 @@ class Consumer:
                 # place. If this fails, we'll also rollback the transaction.
 
                 create_source_path_directory(document.source_path)
-                self._write(document, filename, document.source_path)
+                self._write(document, self.path, document.source_path)
                 self._write(document, thumbnail, document.thumbnail_path)
 
                 # Delete the file only if it was successfully consumed
-                self.log("debug", "Deleting document {}".format(filename))
-                os.unlink(filename)
+                self.log("debug", "Deleting file {}".format(self.path))
+                os.unlink(self.path)
         except Exception as e:
             raise ConsumerError(e)
         finally:
@@ -207,25 +204,25 @@ class Consumer:
 
         return document
 
-    def _store(self, text, doc, thumbnail, date,
-               original_filename=None,
-               force_title=None,
-               force_correspondent_id=None,
-               force_document_type_id=None,
-               force_tag_ids=None):
+    def _store(self, text, date):
 
         # If someone gave us the original filename, use it instead of doc.
 
-        file_info = FileInfo.from_path(original_filename or doc)
+        file_info = FileInfo.from_path(self.filename)
 
-        stats = os.stat(doc)
+        stats = os.stat(self.path)
 
         self.log("debug", "Saving record to database")
 
         created = file_info.created or date or timezone.make_aware(
             datetime.datetime.fromtimestamp(stats.st_mtime))
 
-        with open(doc, "rb") as f:
+        if settings.PASSPHRASE:
+            storage_type = Document.STORAGE_TYPE_GPG
+        else:
+            storage_type = Document.STORAGE_TYPE_UNENCRYPTED
+
+        with open(self.path, "rb") as f:
             document = Document.objects.create(
                 correspondent=file_info.correspondent,
                 title=file_info.title,
@@ -234,7 +231,7 @@ class Consumer:
                 checksum=hashlib.md5(f.read()).hexdigest(),
                 created=created,
                 modified=created,
-                storage_type=self.storage_type
+                storage_type=storage_type
             )
 
         relevant_tags = set(file_info.tags)
@@ -243,18 +240,7 @@ class Consumer:
             self.log("debug", "Tagging with {}".format(tag_names))
             document.tags.add(*relevant_tags)
 
-        if force_title:
-            document.title = force_title
-
-        if force_correspondent_id:
-            document.correspondent = Correspondent.objects.get(pk=force_correspondent_id)
-
-        if force_document_type_id:
-            document.document_type = DocumentType.objects.get(pk=force_document_type_id)
-
-        if force_tag_ids:
-            for tag_id in force_tag_ids:
-                document.tags.add(Tag.objects.get(pk=tag_id))
+        self.apply_overrides(document)
 
         document.filename = generate_filename(document)
 
@@ -264,6 +250,20 @@ class Consumer:
 
         return document
 
+    def apply_overrides(self, document):
+        if self.override_title:
+            document.title = self.override_title
+
+        if self.override_correspondent_id:
+            document.correspondent = Correspondent.objects.get(pk=self.override_correspondent_id)
+
+        if self.override_document_type_id:
+            document.document_type = DocumentType.objects.get(pk=self.override_document_type_id)
+
+        if self.override_tag_ids:
+            for tag_id in self.override_tag_ids:
+                document.tags.add(Tag.objects.get(pk=tag_id))
+
     def _write(self, document, source, target):
         with open(source, "rb") as read_file:
             with open(target, "wb") as write_file: