Coverage for presidio_structured / structured_engine.py: 92%

1import logging 

2from typing import Dict, Optional, Union 

3 

4from pandas import DataFrame 

5from presidio_anonymizer.entities import OperatorConfig 

6 

7from presidio_structured.config import StructuredAnalysis 

8from presidio_structured.data.data_processors import ( 

9 DataProcessorBase, 

10 PandasDataProcessor, 

11) 

12 

13DEFAULT = "replace" 

14 

15 

16class StructuredEngine: 

17 """Class to implement methods for anonymizing tabular data.""" 

18 

19 def __init__(self, data_processor: Optional[DataProcessorBase] = None) -> None: 

20 """ 

21 Initialize the class with a data processor. 

22 

23 :param data_processor: Instance of DataProcessorBase. 

24 """ 

25 if data_processor is None: 

26 self.data_processor = PandasDataProcessor() 

27 else: 

28 self.data_processor = data_processor 

29 

30 self.logger = logging.getLogger("presidio-structured") 

31 

32 def anonymize( 

33 self, 

34 data: Union[Dict, DataFrame], 

35 structured_analysis: StructuredAnalysis, 

36 operators: Union[Dict[str, OperatorConfig], None] = None, 

37 ) -> Union[Dict, DataFrame]: 

38 """ 

39 Anonymize the given data using the given configuration. 

40 

41 :param data: input data as dictionary or pandas DataFrame. 

42 :param structured_analysis: structured analysis configuration. 

43 :param operators: a dictionary of operator configurations, optional. 

44 :return: Anonymized dictionary or DataFrame. 

45 """ 

46 self.logger.debug("Starting anonymization") 

47 operators = self.__check_or_add_default_operator(operators) 

48 

49 return self.data_processor.operate(data, structured_analysis, operators) 

50 

51 def __check_or_add_default_operator( 

52 self, operators: Union[Dict[str, OperatorConfig], None] 

53 ) -> Dict[str, OperatorConfig]: 

54 """ 

55 Check if the provided operators dictionary has a default operator. If not, add a default operator. 

56 

57 :param operators: dictionary of operator configurations. 

58 :return: operators dictionary with the default operator added \ 

59 if it was not initially present. 

60 """ # noqa: E501 

61 default_operator = OperatorConfig(DEFAULT) 

62 if not operators: 

63 self.logger.debug("No operators provided, using default operator") 

64 return {"DEFAULT": default_operator} 

65 if not operators.get("DEFAULT"): 

66 self.logger.debug("No default operator provided, using default operator") 

67 operators["DEFAULT"] = default_operator 

68 return operators