Adding more typing around the classification and matching

This commit is contained in:
Trenton Holmes
2023-07-23 16:49:20 -07:00
committed by Trenton H
parent 07e7bcd30b
commit d376f9e7a3
3 changed files with 30 additions and 24 deletions

View File

@@ -1,7 +1,9 @@
import logging
import re
from documents.classifier import DocumentClassifier
from documents.models import Correspondent
from documents.models import Document
from documents.models import DocumentType
from documents.models import MatchingModel
from documents.models import StoragePath
@@ -11,7 +13,7 @@ from documents.permissions import get_objects_for_user_owner_aware
logger = logging.getLogger("paperless.matching")
def log_reason(matching_model, document, reason):
def log_reason(matching_model: MatchingModel, document: Document, reason: str):
class_name = type(matching_model).__name__
logger.debug(
f"{class_name} {matching_model.name} matched on document "
@@ -19,7 +21,7 @@ def log_reason(matching_model, document, reason):
)
def match_correspondents(document, classifier, user=None):
def match_correspondents(document: Document, classifier: DocumentClassifier, user=None):
pred_id = classifier.predict_correspondent(document.content) if classifier else None
if user is None and document.owner is not None:
@@ -43,7 +45,7 @@ def match_correspondents(document, classifier, user=None):
)
def match_document_types(document, classifier, user=None):
def match_document_types(document: Document, classifier: DocumentClassifier, user=None):
pred_id = classifier.predict_document_type(document.content) if classifier else None
if user is None and document.owner is not None:
@@ -67,7 +69,7 @@ def match_document_types(document, classifier, user=None):
)
def match_tags(document, classifier, user=None):
def match_tags(document: Document, classifier: DocumentClassifier, user=None):
predicted_tag_ids = classifier.predict_tags(document.content) if classifier else []
if user is None and document.owner is not None:
@@ -90,7 +92,7 @@ def match_tags(document, classifier, user=None):
)
def match_storage_paths(document, classifier, user=None):
def match_storage_paths(document: Document, classifier: DocumentClassifier, user=None):
pred_id = classifier.predict_storage_path(document.content) if classifier else None
if user is None and document.owner is not None:
@@ -114,7 +116,7 @@ def match_storage_paths(document, classifier, user=None):
)
def matches(matching_model, document):
def matches(matching_model: MatchingModel, document: Document):
search_kwargs = {}
document_content = document.content