diff options
author | Baruch Sterin <baruchs@gmail.com> | 2011-01-13 22:12:37 +0200 |
---|---|---|
committer | Baruch Sterin <baruchs@gmail.com> | 2011-01-13 22:12:37 +0200 |
commit | 811f5631a812968ccdbe157549f2445747053d50 (patch) | |
tree | d86691b0af7681aba69935875e5a7d779e24bc36 /src/python/redirect.py | |
parent | c85a763444f696e0125445b5adc675020d736c5d (diff) | |
download | abc-811f5631a812968ccdbe157549f2445747053d50.tar.gz abc-811f5631a812968ccdbe157549f2445747053d50.tar.bz2 abc-811f5631a812968ccdbe157549f2445747053d50.zip |
pyabc: reorganize supporting python scripts
Diffstat (limited to 'src/python/redirect.py')
-rw-r--r-- | src/python/redirect.py | 94 |
1 files changed, 94 insertions, 0 deletions
diff --git a/src/python/redirect.py b/src/python/redirect.py new file mode 100644 index 00000000..498fe150 --- /dev/null +++ b/src/python/redirect.py @@ -0,0 +1,94 @@ +""" + +A simple context manager for redirecting streams in Python. +The streams are redirected at the the C runtime level so that the output of C extensions +that use stdio will also be redirected. + +null_file : a stream representing the null device (e.g. /dev/null on Unix) +redirect: a context manager for redirecting streams + +Author: Baruch Sterin (sterin@berkeley.edu) + +""" + +import os +import sys + +from contextlib import contextmanager + +null_file = open( os.devnull, "w" ) + +@contextmanager +def _dup( f ): + fd = os.dup( f.fileno() ) + yield fd + os.close(fd) + +@contextmanager +def redirect(dst = null_file, src = sys.stdout): + + """ + Redirect the src stream into dst. + + Example: + with redirect( open("somefile.txt", sys.stdout ) ): + do some stuff ... + """ + + if src.fileno() == dst.fileno(): + yield + return + + with _dup( src ) as fd_dup_src: + + dst.flush() + + src.flush() + os.close( src.fileno() ) + os.dup2( dst.fileno(), src.fileno() ) + + yield + + src.flush() + os.close( src.fileno() ) + os.dup2( fd_dup_src, src.fileno() ) + +def start_redirect(dst = null_file, src = sys.stdout): + + """ + Start redirection of src stream into dst. Return the duplicated file handle of the source. + + Example: + fd = start_redirect( open("somefile.txt"), sys.stdout ) + ... do some stuff ... + end_redirect(sys.stdout, fd) + """ + + if src.fileno() == dst.fileno(): + return None + + fd_dup_src = os.dup( src.fileno() ) + + dst.flush() + src.flush() + + os.close( src.fileno() ) + os.dup2( dst.fileno(), src.fileno() ) + + return fd_dup_src + +def end_redirect(src, fd_dup_src): + + """ + End redirection of stream src.Redirect the src stream into dst. src is the source stream and fd_dup_src is the value returned by + start_redirect() + """ + + if fd_dup_src is None: + return + + src.flush() + os.close( src.fileno() ) + os.dup2( fd_dup_src, src.fileno() ) + + os.close(fd_dup_src) |