1# -*- coding: utf-8 -*-
2# Copyright (C) 2019-2021 Greenbone Networks GmbH
3#
4# SPDX-License-Identifier: GPL-3.0-or-later
5#
6# This program is free software: you can redistribute it and/or modify
7# it under the terms of the GNU General Public License as published by
8# the Free Software Foundation, either version 3 of the License, or
9# (at your option) any later version.
10#
11# This program is distributed in the hope that it will be useful,
12# but WITHOUT ANY WARRANTY; without even the implied warranty of
13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14# GNU General Public License for more details.
15#
16# You should have received a copy of the GNU General Public License
17# along with this program.  If not, see <http://www.gnu.org/licenses/>.
18
19import sys
20import os
21
22
23class SuppressOutput:
24    def __init__(self, *, suppress_stdout=False, suppress_stderr=False):
25        self.suppress_stdout = suppress_stdout
26        self.suppress_stderr = suppress_stderr
27        self.original_stdout = None
28        self.original_stderr = None
29
30    def __enter__(self):
31        self.devnull = open(os.devnull, "w", encoding='utf-8')
32
33        # Suppress streams
34        if self.suppress_stdout:
35            self.original_stdout = sys.stdout
36            sys.stdout = self.devnull
37
38        if self.suppress_stderr:
39            self.original_stderr = sys.stderr
40            sys.stderr = self.devnull
41
42    def __exit__(self, *args, **kwargs):
43        # Restore streams
44        if self.suppress_stdout:
45            sys.stdout = self.original_stdout
46
47        if self.suppress_stderr:
48            sys.stderr = self.original_stderr
49
50        self.devnull.close()
51