Coverage for presidio_cli / analyzer.py: 98%

1from typing import Generator, Optional, Union 

2 

3from presidio_analyzer import RecognizerResult 

4 

5from presidio_cli.config import PresidioCLIConfig 

6 

7 

8class Line(object): 

9 """Represents a line of text source.""" 

10 

11 def __init__(self, line_no: int, buffer: str, start: int, end: int) -> None: 

12 self.line_no = line_no 

13 self.start = start 

14 self.end = end 

15 self.buffer = buffer 

16 

17 @property 

18 def content(self): 

19 """ 

20 Get the content of the line. 

21 

22 :returns: The encrypted text. 

23 """ 

24 

25 return self.buffer[self.start : self.end] 

26 

27 

28def line_generator(buffer: str) -> Generator[Line, None, None]: 

29 """Generate Line objects from text source. Returns a generator of Line objects. 

30 

31 :param buffer: str, string to read from 

32 """ 

33 line_no = 1 

34 cur = 0 

35 next = buffer.find("\n") 

36 while next != -1: 

37 if next > 0 and buffer[next - 1] == "\r": 

38 yield Line(line_no, buffer, start=cur, end=next - 1) 

39 else: 

40 yield Line(line_no, buffer, start=cur, end=next) 

41 cur = next + 1 

42 next = buffer.find("\n", cur) 

43 line_no += 1 

44 

45 yield Line(line_no, buffer, start=cur, end=len(buffer)) 

46 

47 

48class PIIProblem(object): 

49 """Represents a PII problem found by presidio-cli.""" 

50 

51 def __init__(self, line: int, recognizer_result: RecognizerResult) -> None: 

52 assert isinstance(recognizer_result, RecognizerResult) 

53 self.recognizer_result = recognizer_result.to_dict() 

54 #: Line on which the problem was found (starting at 1) 

55 self.line = line 

56 #: Column on which the problem was found (starting at 1) 

57 self.column = self.recognizer_result["start"] + 1 

58 #: Human-readable description of the problem 

59 self.explanation = self.recognizer_result["analysis_explanation"] 

60 #: Identifier of the rule that detected the problem 

61 self.type = self.recognizer_result["entity_type"] 

62 # Score as a probability determined by the model 

63 self.score = self.recognizer_result["score"] 

64 

65 

66def _analyze( 

67 buffer: str, conf: "PresidioCLIConfig" 

68) -> Generator["PIIProblem", None, None]: 

69 """Analyze a text source. Returns a generator of PIIProblem objects. 

70 

71 :param buffer: str, string to read from 

72 :param conf: presidio_cli configuration object 

73 """ 

74 assert hasattr( 

75 buffer, "__getitem__" 

76 ), "_run() argument must be a buffer, not a stream" 

77 

78 for line in line_generator(buffer): 

79 for result in conf.analyzer.analyze( 

80 text=line.content, 

81 entities=conf.entities, 

82 language=conf.language, 

83 allow_list=conf.allow_list, 

84 ): 

85 p = PIIProblem(line.line_no, result) 

86 if p.score >= conf.threshold: 

87 yield p 

88 

89 

90def analyze( 

91 input: str, conf: "PresidioCLIConfig", filepath: Optional[str] = None 

92) -> Union[tuple, Generator["PIIProblem", None, None], None]: 

93 """Analyze a text source. Returns a generator of PIIProblem objects. 

94 

95 :param input: buffer, string or stream to read from 

96 :param conf: presidio_cli configuration object 

97 :param filepath: string, string with path to file 

98 """ 

99 

100 if conf.is_file_ignored(filepath): 

101 return tuple() 

102 if isinstance(input, (bytes, str)): 

103 return _analyze(input, conf) 

104 elif hasattr(input, "read"): # Python 2's file or Python 3's io.IOBase 

105 # We need to have everything in memory to parse correctly 

106 content = input.read() 

107 return _analyze(content, conf) 

108 else: 

109 raise TypeError("input should be a string or a stream")