Coverage for paperless_ai/matching.py: 39%

1import difflib 

2import logging 

3import re 

4from typing import TypeVar 

5 

6from django.contrib.auth.models import User 

7from django.db.models import Model 

8from django.db.models import QuerySet 

9 

10from documents.models import Correspondent 

11from documents.models import DocumentType 

12from documents.models import StoragePath 

13from documents.models import Tag 

14from documents.permissions import permitted_object_ids 

15from documents.permissions import restrict_queryset_to_visible 

16 

17MATCH_THRESHOLD = 0.8 

18 

19logger = logging.getLogger("paperless_ai.matching") 

20 

21ModelT = TypeVar("ModelT", bound=Model) 

22 

23 

24def _resolve_visible_ids( 

25 ids: list[int], 

26 user: User | None, 

27 model: type[ModelT], 

28 perm: str, 

29) -> list[ModelT]: 

30 """Resolve model-returned IDs against what the user may currently see. 

31 Invalid, deleted, or now-invisible IDs are silently dropped - the model's 

32 belief that an ID exists and is visible may be stale by the time the 

33 response comes back. 

34 """ 

35 if not ids: 

36 return [] 

37 queryset = restrict_queryset_to_visible( 

38 model.objects.filter(pk__in=ids), 

39 user, 

40 perm, 

41 ) 

42 return list(queryset) 

43 

44 

45def resolve_tag_ids(ids: list[int], user: User | None) -> list[Tag]: 

46 return _resolve_visible_ids(ids, user, Tag, "view_tag") 

47 

48 

49def resolve_correspondent_ids( 

50 ids: list[int], 

51 user: User | None, 

52) -> list[Correspondent]: 

53 return _resolve_visible_ids(ids, user, Correspondent, "view_correspondent") 

54 

55 

56def resolve_document_type_ids(ids: list[int], user: User | None) -> list[DocumentType]: 

57 return _resolve_visible_ids(ids, user, DocumentType, "view_documenttype") 

58 

59 

60def resolve_storage_path_ids(ids: list[int], user: User | None) -> list[StoragePath]: 

61 return _resolve_visible_ids(ids, user, StoragePath, "view_storagepath") 

62 

63 

64def _match_by_name( 

65 names: list[str], 

66 user: User | None, 

67 model: type[ModelT], 

68 perm: str, 

69) -> list[ModelT]: 

70 # A workflow may have no user. In that case permitted_object_ids limits 

71 # matching to unowned objects, avoiding another user's private taxonomy. 

72 queryset = model.objects.filter( 

73 pk__in=permitted_object_ids(user, model, perm), 

74 ) 

75 return _match_names_to_queryset(names, queryset) 

76 

77 

78def match_tags_by_name(names: list[str], user: User | None) -> list[Tag]: 

79 return _match_by_name(names, user, Tag, "view_tag") 

80 

81 

82def match_correspondents_by_name( 

83 names: list[str], 

84 user: User | None, 

85) -> list[Correspondent]: 

86 return _match_by_name(names, user, Correspondent, "view_correspondent") 

87 

88 

89def match_document_types_by_name( 

90 names: list[str], 

91 user: User | None, 

92) -> list[DocumentType]: 

93 return _match_by_name(names, user, DocumentType, "view_documenttype") 

94 

95 

96def match_storage_paths_by_name( 

97 names: list[str], 

98 user: User | None, 

99) -> list[StoragePath]: 

100 return _match_by_name(names, user, StoragePath, "view_storagepath") 

101 

102 

103def _normalize(s: str) -> str: 

104 s = s.lower() 

105 s = re.sub(r"[^\w\s]", "", s) # remove punctuation 

106 s = s.strip() 

107 return s 

108 

109 

110def _match_names_to_queryset( 

111 names: list[str], 

112 queryset: QuerySet[ModelT], 

113 attr: str = "name", 

114) -> list[ModelT]: 

115 """Match each name to at most one object, exactly first and fuzzily as a 

116 fallback. A matched object is removed from the pool so two names can never 

117 resolve to the same object; names that match nothing are simply skipped. 

118 """ 

119 results: list[ModelT] = [] 

120 objects = list(queryset) 

121 object_names = [_normalize(getattr(obj, attr)) for obj in objects] 

122 

123 for name in names: 

124 if not name: 

125 continue 

126 target = _normalize(name) 

127 

128 if target in object_names: 

129 index = object_names.index(target) 

130 else: 

131 matches = difflib.get_close_matches( 

132 target, 

133 object_names, 

134 n=1, 

135 cutoff=MATCH_THRESHOLD, 

136 ) 

137 if not matches: 

138 continue 

139 index = object_names.index(matches[0]) 

140 

141 object_names.pop(index) # keep both lists aligned after removal 

142 results.append(objects.pop(index)) 

143 return results 

144 

145 

146def extract_unmatched_names( 

147 names: list[str], 

148 matched_objects: list, 

149 attr="name", 

150) -> list[str]: 

151 matched_names = {_normalize(getattr(obj, attr)) for obj in matched_objects} 

152 return [name for name in names if _normalize(name) not in matched_names]