Coverage for presidio_analyzer / batch_analyzer_engine.py: 98%

1import logging 

2from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union 

3 

4from presidio_analyzer import AnalyzerEngine, DictAnalyzerResult, RecognizerResult 

5from presidio_analyzer.nlp_engine import NlpArtifacts 

6 

7logger = logging.getLogger("presidio-analyzer") 

8 

9 

10class BatchAnalyzerEngine: 

11 """ 

12 Batch analysis of documents (tables, lists, dicts). 

13 

14 Wrapper class to run Presidio Analyzer Engine on multiple values, 

15 either lists/iterators of strings, or dictionaries. 

16 

17 :param analyzer_engine: AnalyzerEngine instance to use 

18 for handling the values in those collections. 

19 """ 

20 

21 def __init__(self, analyzer_engine: Optional[AnalyzerEngine] = None): 

22 self.analyzer_engine = analyzer_engine 

23 if not analyzer_engine: 

24 self.analyzer_engine = AnalyzerEngine() 

25 

26 def analyze_iterator( 

27 self, 

28 texts: Iterable[Union[str, bool, float, int]], 

29 language: str, 

30 batch_size: int = 1, 

31 n_process: int = 1, 

32 **kwargs, 

33 ) -> List[List[RecognizerResult]]: 

34 """ 

35 Analyze an iterable of strings. 

36 

37 :param texts: An list containing strings to be analyzed. 

38 :param language: Input language 

39 :param batch_size: Batch size to process in a single iteration 

40 :param n_process: Number of processors to use. Defaults to `1` 

41 :param kwargs: Additional parameters for the `AnalyzerEngine.analyze` method. 

42 (default value depends on the nlp engine implementation) 

43 """ 

44 

45 # validate types 

46 texts = self._validate_types(texts) 

47 

48 # Process the texts as batch for improved performance 

49 nlp_artifacts_batch: Iterator[Tuple[str, NlpArtifacts]] = ( 

50 self.analyzer_engine.nlp_engine.process_batch( 

51 texts=texts, 

52 language=language, 

53 batch_size=batch_size, 

54 n_process=n_process, 

55 ) 

56 ) 

57 

58 list_results = [] 

59 for text, nlp_artifacts in nlp_artifacts_batch: 

60 results = self.analyzer_engine.analyze( 

61 text=str(text), nlp_artifacts=nlp_artifacts, language=language, **kwargs 

62 ) 

63 

64 list_results.append(results) 

65 

66 return list_results 

67 

68 def analyze_dict( 

69 self, 

70 input_dict: Dict[str, Union[Any, Iterable[Any]]], 

71 language: str, 

72 keys_to_skip: Optional[List[str]] = None, 

73 batch_size: int = 1, 

74 n_process: int = 1, 

75 **kwargs, 

76 ) -> Iterator[DictAnalyzerResult]: 

77 """ 

78 Analyze a dictionary of keys (strings) and values/iterable of values. 

79 

80 Non-string values are returned as is. 

81 

82 :param input_dict: The input dictionary for analysis 

83 :param language: Input language 

84 :param keys_to_skip: Keys to ignore during analysis 

85 :param batch_size: Batch size to process in a single iteration 

86 :param n_process: Number of processors to use. Defaults to `1` 

87 

88 :param kwargs: Additional keyword arguments 

89 for the `AnalyzerEngine.analyze` method. 

90 Use this to pass arguments to the analyze method, 

91 such as `ad_hoc_recognizers`, `context`, `return_decision_process`. 

92 See `AnalyzerEngine.analyze` for the full list. 

93 """ 

94 

95 context = [] 

96 if "context" in kwargs: 

97 context = kwargs["context"] 

98 del kwargs["context"] 

99 

100 if not keys_to_skip: 

101 keys_to_skip = [] 

102 

103 for key, value in input_dict.items(): 

104 if not value or key in keys_to_skip: 

105 yield DictAnalyzerResult(key=key, value=value, recognizer_results=[]) 

106 continue # skip this key as requested 

107 

108 # Add the key as an additional context 

109 specific_context = context[:] 

110 specific_context.append(key) 

111 

112 if type(value) in (str, int, bool, float): 

113 results: List[RecognizerResult] = self.analyzer_engine.analyze( 

114 text=str(value), language=language, context=[key], **kwargs 

115 ) 

116 elif isinstance(value, dict): 

117 new_keys_to_skip = self._get_nested_keys_to_skip(key, keys_to_skip) 

118 results = self.analyze_dict( 

119 input_dict=value, 

120 language=language, 

121 context=specific_context, 

122 keys_to_skip=new_keys_to_skip, 

123 **kwargs, 

124 ) 

125 elif isinstance(value, Iterable): 

126 # Recursively iterate nested dicts 

127 

128 results: List[List[RecognizerResult]] = self.analyze_iterator( 

129 texts=value, 

130 language=language, 

131 context=specific_context, 

132 n_process=n_process, 

133 batch_size=batch_size, 

134 **kwargs, 

135 ) 

136 else: 

137 raise ValueError(f"type {type(value)} is unsupported.") 

138 

139 yield DictAnalyzerResult(key=key, value=value, recognizer_results=results) 

140 

141 @staticmethod 

142 def _validate_types(value_iterator: Iterable[Any]) -> Iterator[Any]: 

143 for val in value_iterator: 

144 if val and type(val) not in (int, float, bool, str): 

145 err_msg = ( 

146 "Analyzer.analyze_iterator only works " 

147 "on primitive types (int, float, bool, str). " 

148 "Lists of objects are not yet supported." 

149 ) 

150 logger.error(err_msg) 

151 raise ValueError(err_msg) 

152 yield val 

153 

154 @staticmethod 

155 def _get_nested_keys_to_skip(key, keys_to_skip): 

156 new_keys_to_skip = [ 

157 k.replace(f"{key}.", "") for k in keys_to_skip if k.startswith(key) 

158 ] 

159 return new_keys_to_skip