Coverage for presidio_structured / analysis_builder.py: 89%
119 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-03-29 08:58 +0000
1import logging
2from abc import ABC, abstractmethod
3from collections import Counter
4from collections.abc import Iterable
5from typing import Dict, Iterator, List, Optional, Union
7from pandas import DataFrame
8from presidio_analyzer import (
9 AnalyzerEngine,
10 BatchAnalyzerEngine,
11 DictAnalyzerResult,
12 RecognizerResult,
13)
15from presidio_structured.config import StructuredAnalysis
17NON_PII_ENTITY_TYPE = "NON_PII"
19logger = logging.getLogger("presidio-structured")
22class AnalysisBuilder(ABC):
23 """Abstract base class for a configuration generator."""
25 def __init__(
26 self,
27 analyzer: Optional[AnalyzerEngine] = None,
28 analyzer_score_threshold: Optional[float] = None,
29 n_process: int = 1,
30 batch_size: int = 1
31 ) -> None:
32 """Initialize the configuration generator.
34 :param analyzer: AnalyzerEngine instance
35 :param analyzer_score_threshold: threshold for filtering out results
36 :param batch_size: Batch size to process in a single iteration
37 :param n_process: Number of processors to use. Defaults to `1`
38 """
39 default_score_threshold = (
40 analyzer_score_threshold if analyzer_score_threshold is not None else 0
41 )
42 self.analyzer = (
43 AnalyzerEngine(default_score_threshold=default_score_threshold)
44 if analyzer is None
45 else analyzer
46 )
47 self.batch_analyzer = BatchAnalyzerEngine(analyzer_engine=self.analyzer)
48 self.n_process = n_process
49 self.batch_size = batch_size
51 @abstractmethod
52 def generate_analysis(
53 self,
54 data: Union[Dict, DataFrame],
55 language: str = "en",
56 score_threshold: Optional[float] = None,
57 ) -> StructuredAnalysis:
58 """
59 Abstract method to generate a configuration from the given data.
61 :param data: The input data. Can be a dictionary or DataFrame instance.
62 :return: The generated configuration.
63 """
64 pass
66 def _remove_low_scores(
67 self,
68 key_recognizer_result_map: Dict[str, RecognizerResult],
69 score_threshold: Optional[float] = None,
70 ) -> Dict[str, RecognizerResult]:
71 """
72 Remove results for which the confidence is lower than the threshold.
74 :param results: Dict of column names to RecognizerResult
75 :param score_threshold: float value for minimum possible confidence
76 :return: List[RecognizerResult]
77 """
78 if score_threshold is None:
79 score_threshold = self.analyzer.default_score_threshold
81 new_key_recognizer_result_map = {}
82 for column, result in key_recognizer_result_map.items():
83 if result.score >= score_threshold:
84 new_key_recognizer_result_map[column] = result
86 return new_key_recognizer_result_map
89class JsonAnalysisBuilder(AnalysisBuilder):
90 """Concrete configuration generator for JSON data."""
92 def generate_analysis(
93 self,
94 data: Dict,
95 language: str = "en",
96 ) -> StructuredAnalysis:
97 """
98 Generate a configuration from the given JSON data.
100 :param data: The input JSON data.
101 :return: The generated configuration.
102 """
103 logger.debug("Starting JSON BatchAnalyzer analysis")
104 analyzer_results = self.batch_analyzer.analyze_dict(
105 input_dict=data,
106 language=language,
107 n_process=self.n_process,
108 batch_size=self.batch_size
109 )
111 key_recognizer_result_map = self._generate_analysis_from_results_json(
112 analyzer_results
113 )
115 key_entity_map = {
116 key: result.entity_type for key, result in key_recognizer_result_map.items()
117 }
119 return StructuredAnalysis(entity_mapping=key_entity_map)
121 def _generate_analysis_from_results_json(
122 self, analyzer_results: Iterator[DictAnalyzerResult], prefix: str = ""
123 ) -> Dict[str, RecognizerResult]:
124 """
125 Generate a configuration from the given analyzer results. Always uses the first recognizer result if there are more than one.
127 :param analyzer_results: The analyzer results.
128 :param prefix: The prefix for the configuration keys.
129 :return: The generated configuration.
130 """ # noqa: E501
131 key_recognizer_result_map = {}
133 if not isinstance(analyzer_results, Iterable):
134 logger.debug(
135 "No analyzer results found, returning empty StructuredAnalysis"
136 )
137 return key_recognizer_result_map
139 for result in analyzer_results:
140 current_key = prefix + result.key
142 if isinstance(result.value, dict) and isinstance(
143 result.recognizer_results, Iterator
144 ):
145 nested_mappings = self._generate_analysis_from_results_json(
146 result.recognizer_results, prefix=current_key + "."
147 )
148 key_recognizer_result_map.update(nested_mappings)
149 first_recognizer_result = next(iter(result.recognizer_results), None)
150 if isinstance(first_recognizer_result, RecognizerResult):
151 logger.debug(
152 f"Found result with entity {first_recognizer_result.entity_type} \
153 in {current_key}"
154 )
155 key_recognizer_result_map[current_key] = first_recognizer_result
156 return key_recognizer_result_map
159class TabularAnalysisBuilder(AnalysisBuilder):
160 """Placeholder class for generalizing tabular data analysis builders (e.g. PySpark). Only implemented as PandasAnalysisBuilder for now.""" # noqa: E501
162 pass
165class PandasAnalysisBuilder(TabularAnalysisBuilder):
166 """Concrete configuration generator for tabular data."""
168 entity_selection_strategies = {"highest_confidence", "mixed", "most_common"}
170 def generate_analysis(
171 self,
172 df: DataFrame,
173 n: Optional[int] = None,
174 language: str = "en",
175 selection_strategy: str = "most_common",
176 mixed_strategy_threshold: float = 0.5,
177 ) -> StructuredAnalysis:
178 """
179 Generate a configuration from the given tabular data.
181 :param df: The input tabular data (dataframe).
182 :param n: The number of samples to be taken from the dataframe.
183 :param language: The language to be used for analysis.
184 :param selection_strategy: A string that specifies the entity selection strategy
185 ('highest_confidence', 'mixed', or default to most common).
186 :param mixed_strategy_threshold: A float value for the threshold to be used in
187 the entity selection mixed strategy.
188 :return: A StructuredAnalysis object containing the analysis results.
189 """
190 if not n:
191 n = len(df)
192 elif n > len(df):
193 logger.debug(
194 f"Number of samples ({n}) is larger than the number of rows \
195 ({len(df)}), using all rows"
196 )
197 n = len(df)
199 df = df.sample(n, random_state=123)
201 key_recognizer_result_map = self._generate_key_rec_results_map(
202 df, language, selection_strategy, mixed_strategy_threshold
203 )
205 key_entity_map = {
206 key: result.entity_type
207 for key, result in key_recognizer_result_map.items()
208 if result.entity_type != NON_PII_ENTITY_TYPE
209 }
211 return StructuredAnalysis(entity_mapping=key_entity_map)
213 def _generate_key_rec_results_map(
214 self,
215 df: DataFrame,
216 language: str,
217 selection_strategy: str = "most_common",
218 mixed_strategy_threshold: float = 0.5,
219 ) -> Dict[str, RecognizerResult]:
220 """
221 Find the most common entity in a dataframe column.
223 If more than one entity is found in a cell, the first one is used.
225 :param df: The dataframe where entities will be searched.
226 :param language: Language to be used in the analysis engine.
227 :param selection_strategy: A string that specifies the entity selection strategy
228 ('highest_confidence', 'mixed', or default to most common).
229 :param mixed_strategy_threshold: A float value for the threshold to be used in
230 the entity selection mixed strategy.
231 :return: A dictionary mapping column names to the most common RecognizerResult.
232 """
233 column_analyzer_results_map = self._batch_analyze_df(df, language)
234 key_recognizer_result_map = {}
235 for column, analyzer_result in column_analyzer_results_map.items():
236 key_recognizer_result_map[column] = self._find_entity_based_on_strategy(
237 analyzer_result, selection_strategy, mixed_strategy_threshold
238 )
239 return key_recognizer_result_map
241 def _batch_analyze_df(
242 self, df: DataFrame, language: str
243 ) -> Dict[str, List[List[RecognizerResult]]]:
244 """
245 Analyze each column in the dataframe for entities using the batch analyzer.
247 :param df: The dataframe to be analyzed.
248 :param language: The language configuration for the analyzer.
249 :return: A dictionary mapping each column name to a \
250 list of lists of RecognizerResults.
251 """
252 column_analyzer_results_map = {}
253 for column in df.columns:
254 logger.debug(f"Finding most common PII entity for column {column}")
255 analyzer_results = self.batch_analyzer.analyze_iterator(
256 [val for val in df[column]],
257 language=language,
258 n_process=self.n_process,
259 batch_size=self.batch_size
260 )
261 column_analyzer_results_map[column] = analyzer_results
263 return column_analyzer_results_map
265 def _find_entity_based_on_strategy(
266 self,
267 analyzer_results: List[List[RecognizerResult]],
268 selection_strategy: str,
269 mixed_strategy_threshold: float,
270 ) -> RecognizerResult:
271 """
272 Determine the most suitable entity based on the specified selection strategy.
274 :param analyzer_results: A nested list of RecognizerResult objects from the
275 analysis results.
276 :param selection_strategy: A string that specifies the entity selection strategy
277 ('highest_confidence', 'mixed', or default to most common).
278 :return: A RecognizerResult object representing the selected entity based on the
279 given strategy.
280 """
281 if selection_strategy not in self.entity_selection_strategies:
282 raise ValueError(
283 f"Unsupported entity selection strategy: {selection_strategy}."
284 )
286 if not any(analyzer_results):
287 return RecognizerResult(
288 entity_type=NON_PII_ENTITY_TYPE, start=0, end=1, score=1.0
289 )
291 flat_results = self._flatten_results(analyzer_results)
293 # Select the entity based on the desired strategy
294 if selection_strategy == "highest_confidence":
295 return self._select_highest_confidence_entity(flat_results)
296 elif selection_strategy == "mixed":
297 return self._select_mixed_strategy_entity(
298 flat_results, mixed_strategy_threshold
299 )
301 return self._select_most_common_entity(flat_results)
303 def _select_most_common_entity(self, flat_results):
304 """
305 Select the most common entity from the flattened analysis results.
307 :param flat_results: A list of tuples containing index and RecognizerResult
308 objects from the flattened analysis results.
309 :return: A RecognizerResult object for the most commonly found entity type.
310 """
311 # Count occurrences of each entity type
312 type_counter = Counter(res.entity_type for _, res in flat_results)
313 most_common_type, most_common_count = type_counter.most_common(1)[0]
315 # Calculate the score as the proportion of occurrences
316 score = most_common_count / len(flat_results)
318 return RecognizerResult(
319 entity_type=most_common_type, start=0, end=1, score=score
320 )
322 def _select_highest_confidence_entity(self, flat_results):
323 """
324 Select the entity with the highest confidence score.
326 :param flat_results: A list of tuples containing index and RecognizerResult
327 objects from the flattened analysis results.
328 :return: A RecognizerResult object for the entity with the highest confidence
329 score.
330 """
331 score_aggregator = self._aggregate_scores(flat_results)
333 # Find the highest score across all entities
334 highest_score = max(
335 max(scores) for scores in score_aggregator.values() if scores
336 )
338 # Find the entities with the highest score and count their occurrences
339 entities_highest_score = {
340 entity: scores.count(highest_score)
341 for entity, scores in score_aggregator.items()
342 if highest_score in scores
343 }
345 # Find the entity(ies) with the most number of high scores
346 max_occurrences = max(entities_highest_score.values())
347 highest_confidence_entities = [
348 entity
349 for entity, count in entities_highest_score.items()
350 if count == max_occurrences
351 ]
353 return RecognizerResult(
354 entity_type=highest_confidence_entities[0],
355 start=0,
356 end=1,
357 score=highest_score,
358 )
360 def _select_mixed_strategy_entity(self, flat_results, mixed_strategy_threshold):
361 """
362 Select an entity using a mixed strategy.
364 Chooses an entity based on the highest confidence score if it is above the
365 threshold. Otherwise, it defaults to the most common entity.
367 :param flat_results: A list of tuples containing index and RecognizerResult
368 objects from the flattened analysis results.
369 :return: A RecognizerResult object selected based on the mixed strategy.
370 """
371 # Check if mixed strategy threshold is within the valid range
372 if not 0 <= mixed_strategy_threshold <= 1:
373 raise ValueError(
374 f"Invalid mixed strategy threshold: {mixed_strategy_threshold}."
375 )
377 score_aggregator = self._aggregate_scores(flat_results)
379 # Check if the highest score is greater than threshold and select accordingly
380 highest_score = max(
381 max(scores) for scores in score_aggregator.values() if scores
382 )
383 if highest_score > mixed_strategy_threshold:
384 return self._select_highest_confidence_entity(flat_results)
385 else:
386 return self._select_most_common_entity(flat_results)
388 @staticmethod
389 def _aggregate_scores(flat_results):
390 """
391 Aggregate the scores for each entity type from the flattened analysis results.
393 :param flat_results: A list of tuples containing index and RecognizerResult
394 objects from the flattened analysis results.
395 :return: A dictionary with entity types as keys and lists of scores as values.
396 """
397 score_aggregator = {}
398 for _, res in flat_results:
399 if res.entity_type not in score_aggregator:
400 score_aggregator[res.entity_type] = []
401 score_aggregator[res.entity_type].append(res.score)
402 return score_aggregator
404 @staticmethod
405 def _flatten_results(analyzer_results):
406 """
407 Flattens a nested lists of RecognizerResult objects into a list of tuples.
409 :param analyzer_results: A nested list of RecognizerResult objects from
410 the analysis results.
411 :return: A flattened list of tuples containing index and RecognizerResult
412 objects.
413 """
414 return [
415 (cell_idx, res)
416 for cell_idx, cell_results in enumerate(analyzer_results)
417 for res in cell_results
418 ]