from __future__ import print_function
import os
import io
import sys
import traceback
from fcntl import fcntl, F_GETFL, F_SETFL

class MuteStdIO: # export

    def __init__(self, stderr='on', stdout='off'):
        self.__stderr = stderr
        self.__stdout = stdout
        # TODO: arguments not fully implemented,
        #       Add support for 'off', 'on', 'stdxxx' (redirection)
        pass

    def __enter__(self):
        if self.__stdout == 'off':
            rfd, wfd = os.pipe()
            flags = fcntl(rfd, F_GETFL)
            fcntl(rfd, F_SETFL, flags | os.O_NONBLOCK)
            self.rfile = io.open(rfd, 'r')
            self.fake_stdout_stream = io.open(wfd, 'w')
            self.real_stdout_fd = os.dup(1)
            #os.close(1)
            os.dup2(wfd, 1)

    def __exit__(self, type, value, traceback):
        if self.__stdout == 'off':
            sys.stdout.flush()
            os.dup2(self.real_stdout_fd, 1)
            if type is not None:
                    #print("-------- Error while stdio was suppressed --------")
                    #traceback.print_stack()
                    #print(traceback)
                    print("-------- Captured output --------")
                    print(*self.rfile.readlines())
                    self.rfile.close()
                    #print('type = ' + str(type))
                    #print('value = ' + str(value))
                    raise type(value)