1.1 --- a/allpy/util.py	Thu May 31 15:23:37 2012 +0400
     1.2 +++ b/allpy/util.py	Fri Jun 01 13:43:02 2012 +0400
     1.3 @@ -2,6 +2,9 @@
     1.4  """
     1.5  import sys
     1.6  import warnings
     1.7 +import os
     1.8 +from tempfile import mkstemp
     1.9 +from StringIO import StringIO
    1.10  
    1.11  def unzip(seq):
    1.12      """The oppozite of zip() builtin."""
    1.13 @@ -44,4 +47,81 @@
    1.14      """Warn about function being deprecated."""
    1.15      warnings.warn(message, DeprecationWarning, stacklevel=2)
    1.16  
    1.17 +class Silence(object):
    1.18 +    """Context manager for use with `with`.
    1.19 +
    1.20 +    Run code in context without any message to the screen, but show all output
    1.21 +    that was there if an error occured. E.g.::
    1.22 +
    1.23 +        with Silence():
    1.24 +            structure = PDBParser().get_structure(name, file)
    1.25 +
    1.26 +    There are two mutually exclusive ways of silencing:
    1.27 +
    1.28 +        1. By replacing OS'es stdout/stderr with tempfile (called `dup`)
    1.29 +        2. By replacing python's sys.stdout/sys.stderr with StringIO (`hide`)
    1.30 +
    1.31 +    For each of ways you may specify values: 'stdout', 'stderr' or 'both'.
    1.32 +    E.g::
    1.33 +
    1.34 +        with Silence(dup="both"):
    1.35 +            check_call(["mucle", alignment])
    1.36 +
    1.37 +    `Silence()` is equivalent to `Silence(hide='stderr')`
    1.38 +    """
    1.39 +
    1.40 +    def __init__(self, hide=None, dup=None):
    1.41 +        if dup is None and hide is None:
    1.42 +            hide = "stderr"
    1.43 +        assert not (dup is not None and hide is not None)
    1.44 +        self.stderr = self.stdout = self.dup_stderr = self.dup_stdout = False
    1.45 +        if dup == "stdout" or dup == "both":
    1.46 +            self.dup_stdout = True
    1.47 +        if dup == "stderr" or dup == "both":
    1.48 +            self.dup_stderr = True
    1.49 +        if hide == "stdout" or hide == "both":
    1.50 +            self.stdout = True
    1.51 +        if hide == "stderr" or hide == "both":
    1.52 +            self.stderr = True
    1.53 +
    1.54 +    def __enter__(self):
    1.55 +        assert not (self.stderr and self.dup_stderr)
    1.56 +        assert not (self.stdout and self.dup_stdout)
    1.57 +        if self.dup_stdout or self.dup_stderr:
    1.58 +            fd, name = mkstemp(prefix="allpy-silent-")
    1.59 +            self.captured_dup = os.fdopen(fd)
    1.60 +            self.captured_dup_name = name
    1.61 +            if self.dup_stdout:
    1.62 +                self.dup_stdout = os.dup(sys.stdout.fileno())
    1.63 +                os.dup2(self.captured_dup.fileno(), sys.stdout.fileno())
    1.64 +            if self.dup_stderr:
    1.65 +                self.dup_stderr = os.dup(sys.stderr.fileno())
    1.66 +                os.dup2(self.captured_dup.fileno(), sys.stderr.fileno())
    1.67 +        if self.stdout or self.stderr:
    1.68 +            self.captured_output = StringIO()
    1.69 +            if self.stdout:
    1.70 +                self.stdout, sys.stdout = sys.stdout, self.captured_output
    1.71 +            if self.stderr:
    1.72 +                self.stderr, sys.stderr = sys.stderr, self.captured_output
    1.73 +
    1.74 +    def __exit__(self, exc_type, exc_value, traceback):
    1.75 +        if self.stdout or self.stderr:
    1.76 +            if self.stdout:
    1.77 +                sys.stdout = self.stdout
    1.78 +            if self.stderr:
    1.79 +                sys.stderr = self.stderr
    1.80 +            if exc_type is not None:
    1.81 +                sys.stderr.write(self.captured_output.getvalue())
    1.82 +        if self.dup_stdout or self.dup_stderr:
    1.83 +            if self.dup_stdout:
    1.84 +                os.dup2(self.dup_stdout, sys.stdout.fileno())
    1.85 +                os.close(self.dup_stdout)
    1.86 +            if self.dup_stderr:
    1.87 +                os.dup2(self.dup_stderr, sys.stderr.fileno())
    1.88 +                os.close(self.dup_stderr)
    1.89 +            if exc_type is not None:
    1.90 +                self.captured_dup.seek(0)
    1.91 +                sys.stderr.write(self.captured_dup.read())
    1.92 +                os.unlink(self.captured_dup_name)
    1.93 +
    1.94  # vim: set et ts=4 sts=4 sw=4:

     2.1 --- /dev/null	Thu Jan 01 00:00:00 1970 +0000
     2.2 +++ b/test/test_silence.py	Fri Jun 01 13:43:02 2012 +0400
     2.3 @@ -0,0 +1,95 @@
     2.4 +"""Test for allpy.util.Silence
     2.5 +
     2.6 +>>> print check_output(["python", __file__], stderr=STDOUT)
     2.7 +Example1
     2.8 +Stdout must be seen (only stderr silenced)
     2.9 +Example2
    2.10 +Example3
    2.11 +Stdout must be seen (error raised)
    2.12 +Example4
    2.13 +Stdout must be seen (only stderr silenced)
    2.14 +Example5
    2.15 +Example6
    2.16 +Stdout must be seen (error raised)
    2.17 +Stderr must be seen too
    2.18 +...
    2.19 +done
    2.20 +<BLANKLINE>
    2.21 +
    2.22 +"""
    2.23 +import sys
    2.24 +from subprocess import call, check_call, Popen, PIPE, STDOUT
    2.25 +from allpy.util import Silence
    2.26 +
    2.27 +def check_output(*args, **kws):
    2.28 +	"""`check_output` is only introduced in 2.7, we need it NOW"""
    2.29 +	p = Popen(stdout=PIPE, *args, **kws)
    2.30 +	out, _ = p.communicate()
    2.31 +	assert p.poll() == 0
    2.32 +	return out
    2.33 +
    2.34 +class Unbuffered(object):
    2.35 +	"""Fix python stdIO buffering by wrapping this around stdout/stderr.
    2.36 +
    2.37 +	Python uses different buffering strategies in interactive and
    2.38 +	non-interactive invocations making this test output look ugly and
    2.39 +	difficult to understand. This class is a crutch for the problem.
    2.40 +	"""
    2.41 +	def __init__(self, stream):
    2.42 +		self.stream = stream
    2.43 +	def write(self, data):
    2.44 +		self.stream.write(data)
    2.45 +		self.stream.flush()
    2.46 +	def __getattr__(self, attr):
    2.47 +		return getattr(self.stream, attr)
    2.48 +
    2.49 +if __name__ == "__main__":
    2.50 +
    2.51 +	# Only make uniform buffering when running as __main__ program.
    2.52 +	# Otherwise, importing this module would cripple interactive interpreter.
    2.53 +	sys.stdout = Unbuffered(sys.stdout)
    2.54 +	sys.stderr = Unbuffered(sys.stderr)
    2.55 +
    2.56 +	print "Example1"
    2.57 +
    2.58 +	with Silence():
    2.59 +		print "Stdout must be seen (only stderr silenced)"
    2.60 +		sys.stderr.write("Stderr must not be seen")
    2.61 +
    2.62 +	print "Example2"
    2.63 +
    2.64 +	with Silence(hide="both"):
    2.65 +		print "Stdout must not be seen (both silenced)"
    2.66 +
    2.67 +	print "Example3"
    2.68 +
    2.69 +	try:
    2.70 +		with Silence(hide="both"):
    2.71 +			print "Stdout must be seen (error raised)"
    2.72 +			raise Exception()
    2.73 +	except Exception:
    2.74 +		pass
    2.75 +
    2.76 +	call(["echo", "Example4"])
    2.77 +
    2.78 +	with Silence(dup="stderr"):
    2.79 +		check_call(["echo", "Stdout must be seen (only stderr silenced)"])
    2.80 +		check_call(["echo", "Stderr must not be seen"], stdout=sys.stderr)
    2.81 +
    2.82 +	call(["echo", "Example5"])
    2.83 +
    2.84 +	with Silence(dup="both"):
    2.85 +		check_call(["echo", "Stdout must not be seen (both silenced)"])
    2.86 +
    2.87 +	call(["echo", "Example6"])
    2.88 +
    2.89 +	try:
    2.90 +		with Silence(dup="both"):
    2.91 +			check_call(["echo", "Stdout must be seen (error raised)"])
    2.92 +			check_call(["echo", "Stderr must be seen too"], stdout=sys.stderr)
    2.93 +			check_call(["false"])
    2.94 +	except Exception:
    2.95 +		pass
    2.96 +
    2.97 +	print "..."
    2.98 +	call(["echo", "done"])