slither.utils.output_capture

 1import io
 2import logging
 3import sys
 4
 5
 6class CapturingStringIO(io.StringIO):
 7    """
 8    I/O implementation which captures output, and optionally mirrors it to the original I/O stream it replaces.
 9    """
10
11    def __init__(self, original_io=None):
12        super().__init__()
13        self.original_io = original_io
14
15    def write(self, s):
16        super().write(s)
17        if self.original_io:
18            self.original_io.write(s)
19
20
21class StandardOutputCapture:
22    """
23    Redirects and captures standard output/errors.
24    """
25
26    original_stdout = None
27    original_stderr = None
28    original_logger_handlers = None
29
30    @staticmethod
31    def enable(block_original: bool = True) -> None:
32        """
33        Redirects stdout and stderr to a capturable StringIO.
34        :param block_original: If True, blocks all output to the original stream. If False, duplicates output.
35        :return: None
36        """
37        # Redirect stdout
38        if StandardOutputCapture.original_stdout is None:
39            StandardOutputCapture.original_stdout = sys.stdout
40            sys.stdout = CapturingStringIO(
41                None if block_original else StandardOutputCapture.original_stdout
42            )
43
44        # Redirect stderr
45        if StandardOutputCapture.original_stderr is None:
46            StandardOutputCapture.original_stderr = sys.stderr
47            sys.stderr = CapturingStringIO(
48                None if block_original else StandardOutputCapture.original_stderr
49            )
50
51            # Backup and swap root logger handlers
52            root_logger = logging.getLogger()
53            StandardOutputCapture.original_logger_handlers = root_logger.handlers
54            root_logger.handlers = [logging.StreamHandler(sys.stderr)]
55
56    @staticmethod
57    def disable() -> None:
58        """
59        Disables redirection of stdout/stderr, if previously enabled.
60        :return: None
61        """
62        # If we have a stdout backup, restore it.
63        if StandardOutputCapture.original_stdout is not None:
64            sys.stdout.close()
65            sys.stdout = StandardOutputCapture.original_stdout
66            StandardOutputCapture.original_stdout = None
67
68        # If we have an stderr backup, restore it.
69        if StandardOutputCapture.original_stderr is not None:
70            sys.stderr.close()
71            sys.stderr = StandardOutputCapture.original_stderr
72            StandardOutputCapture.original_stderr = None
73
74        # Restore our logging handlers
75        if StandardOutputCapture.original_logger_handlers is not None:
76            root_logger = logging.getLogger()
77            root_logger.handlers = StandardOutputCapture.original_logger_handlers
78            StandardOutputCapture.original_logger_handlers = None
79
80    @staticmethod
81    def get_stdout_output() -> str:
82        """
83        Obtains the output from the currently set stdout
84        :return: Returns stdout output as a string
85        """
86        sys.stdout.seek(0)
87        return sys.stdout.read()
88
89    @staticmethod
90    def get_stderr_output() -> str:
91        """
92        Obtains the output from the currently set stderr
93        :return: Returns stderr output as a string
94        """
95        sys.stderr.seek(0)
96        return sys.stderr.read()