123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371 |
- import os
- import sys
- import py
- import tempfile
- try:
- from io import StringIO
- except ImportError:
- from StringIO import StringIO
- if sys.version_info < (3,0):
- class TextIO(StringIO):
- def write(self, data):
- if not isinstance(data, unicode):
- data = unicode(data, getattr(self, '_encoding', 'UTF-8'), 'replace')
- return StringIO.write(self, data)
- else:
- TextIO = StringIO
- try:
- from io import BytesIO
- except ImportError:
- class BytesIO(StringIO):
- def write(self, data):
- if isinstance(data, unicode):
- raise TypeError("not a byte value: %r" %(data,))
- return StringIO.write(self, data)
- patchsysdict = {0: 'stdin', 1: 'stdout', 2: 'stderr'}
- class FDCapture:
- """ Capture IO to/from a given os-level filedescriptor. """
- def __init__(self, targetfd, tmpfile=None, now=True, patchsys=False):
- """ save targetfd descriptor, and open a new
- temporary file there. If no tmpfile is
- specified a tempfile.Tempfile() will be opened
- in text mode.
- """
- self.targetfd = targetfd
- if tmpfile is None and targetfd != 0:
- f = tempfile.TemporaryFile('wb+')
- tmpfile = dupfile(f, encoding="UTF-8")
- f.close()
- self.tmpfile = tmpfile
- self._savefd = os.dup(self.targetfd)
- if patchsys:
- self._oldsys = getattr(sys, patchsysdict[targetfd])
- if now:
- self.start()
- def start(self):
- try:
- os.fstat(self._savefd)
- except OSError:
- raise ValueError("saved filedescriptor not valid, "
- "did you call start() twice?")
- if self.targetfd == 0 and not self.tmpfile:
- fd = os.open(devnullpath, os.O_RDONLY)
- os.dup2(fd, 0)
- os.close(fd)
- if hasattr(self, '_oldsys'):
- setattr(sys, patchsysdict[self.targetfd], DontReadFromInput())
- else:
- os.dup2(self.tmpfile.fileno(), self.targetfd)
- if hasattr(self, '_oldsys'):
- setattr(sys, patchsysdict[self.targetfd], self.tmpfile)
- def done(self):
- """ unpatch and clean up, returns the self.tmpfile (file object)
- """
- os.dup2(self._savefd, self.targetfd)
- os.close(self._savefd)
- if self.targetfd != 0:
- self.tmpfile.seek(0)
- if hasattr(self, '_oldsys'):
- setattr(sys, patchsysdict[self.targetfd], self._oldsys)
- return self.tmpfile
- def writeorg(self, data):
- """ write a string to the original file descriptor
- """
- tempfp = tempfile.TemporaryFile()
- try:
- os.dup2(self._savefd, tempfp.fileno())
- tempfp.write(data)
- finally:
- tempfp.close()
- def dupfile(f, mode=None, buffering=0, raising=False, encoding=None):
- """ return a new open file object that's a duplicate of f
- mode is duplicated if not given, 'buffering' controls
- buffer size (defaulting to no buffering) and 'raising'
- defines whether an exception is raised when an incompatible
- file object is passed in (if raising is False, the file
- object itself will be returned)
- """
- try:
- fd = f.fileno()
- mode = mode or f.mode
- except AttributeError:
- if raising:
- raise
- return f
- newfd = os.dup(fd)
- if sys.version_info >= (3,0):
- if encoding is not None:
- mode = mode.replace("b", "")
- buffering = True
- return os.fdopen(newfd, mode, buffering, encoding, closefd=True)
- else:
- f = os.fdopen(newfd, mode, buffering)
- if encoding is not None:
- return EncodedFile(f, encoding)
- return f
- class EncodedFile(object):
- def __init__(self, _stream, encoding):
- self._stream = _stream
- self.encoding = encoding
- def write(self, obj):
- if isinstance(obj, unicode):
- obj = obj.encode(self.encoding)
- elif isinstance(obj, str):
- pass
- else:
- obj = str(obj)
- self._stream.write(obj)
- def writelines(self, linelist):
- data = ''.join(linelist)
- self.write(data)
- def __getattr__(self, name):
- return getattr(self._stream, name)
- class Capture(object):
- def call(cls, func, *args, **kwargs):
- """ return a (res, out, err) tuple where
- out and err represent the output/error output
- during function execution.
- call the given function with args/kwargs
- and capture output/error during its execution.
- """
- so = cls()
- try:
- res = func(*args, **kwargs)
- finally:
- out, err = so.reset()
- return res, out, err
- call = classmethod(call)
- def reset(self):
- """ reset sys.stdout/stderr and return captured output as strings. """
- if hasattr(self, '_reset'):
- raise ValueError("was already reset")
- self._reset = True
- outfile, errfile = self.done(save=False)
- out, err = "", ""
- if outfile and not outfile.closed:
- out = outfile.read()
- outfile.close()
- if errfile and errfile != outfile and not errfile.closed:
- err = errfile.read()
- errfile.close()
- return out, err
- def suspend(self):
- """ return current snapshot captures, memorize tempfiles. """
- outerr = self.readouterr()
- outfile, errfile = self.done()
- return outerr
- class StdCaptureFD(Capture):
- """ This class allows to capture writes to FD1 and FD2
- and may connect a NULL file to FD0 (and prevent
- reads from sys.stdin). If any of the 0,1,2 file descriptors
- is invalid it will not be captured.
- """
- def __init__(self, out=True, err=True, mixed=False,
- in_=True, patchsys=True, now=True):
- self._options = {
- "out": out,
- "err": err,
- "mixed": mixed,
- "in_": in_,
- "patchsys": patchsys,
- "now": now,
- }
- self._save()
- if now:
- self.startall()
- def _save(self):
- in_ = self._options['in_']
- out = self._options['out']
- err = self._options['err']
- mixed = self._options['mixed']
- patchsys = self._options['patchsys']
- if in_:
- try:
- self.in_ = FDCapture(0, tmpfile=None, now=False,
- patchsys=patchsys)
- except OSError:
- pass
- if out:
- tmpfile = None
- if hasattr(out, 'write'):
- tmpfile = out
- try:
- self.out = FDCapture(1, tmpfile=tmpfile,
- now=False, patchsys=patchsys)
- self._options['out'] = self.out.tmpfile
- except OSError:
- pass
- if err:
- if out and mixed:
- tmpfile = self.out.tmpfile
- elif hasattr(err, 'write'):
- tmpfile = err
- else:
- tmpfile = None
- try:
- self.err = FDCapture(2, tmpfile=tmpfile,
- now=False, patchsys=patchsys)
- self._options['err'] = self.err.tmpfile
- except OSError:
- pass
- def startall(self):
- if hasattr(self, 'in_'):
- self.in_.start()
- if hasattr(self, 'out'):
- self.out.start()
- if hasattr(self, 'err'):
- self.err.start()
- def resume(self):
- """ resume capturing with original temp files. """
- self.startall()
- def done(self, save=True):
- """ return (outfile, errfile) and stop capturing. """
- outfile = errfile = None
- if hasattr(self, 'out') and not self.out.tmpfile.closed:
- outfile = self.out.done()
- if hasattr(self, 'err') and not self.err.tmpfile.closed:
- errfile = self.err.done()
- if hasattr(self, 'in_'):
- tmpfile = self.in_.done()
- if save:
- self._save()
- return outfile, errfile
- def readouterr(self):
- """ return snapshot value of stdout/stderr capturings. """
- if hasattr(self, "out"):
- out = self._readsnapshot(self.out.tmpfile)
- else:
- out = ""
- if hasattr(self, "err"):
- err = self._readsnapshot(self.err.tmpfile)
- else:
- err = ""
- return out, err
- def _readsnapshot(self, f):
- f.seek(0)
- res = f.read()
- enc = getattr(f, "encoding", None)
- if enc:
- res = py.builtin._totext(res, enc, "replace")
- f.truncate(0)
- f.seek(0)
- return res
- class StdCapture(Capture):
- """ This class allows to capture writes to sys.stdout|stderr "in-memory"
- and will raise errors on tries to read from sys.stdin. It only
- modifies sys.stdout|stderr|stdin attributes and does not
- touch underlying File Descriptors (use StdCaptureFD for that).
- """
- def __init__(self, out=True, err=True, in_=True, mixed=False, now=True):
- self._oldout = sys.stdout
- self._olderr = sys.stderr
- self._oldin = sys.stdin
- if out and not hasattr(out, 'file'):
- out = TextIO()
- self.out = out
- if err:
- if mixed:
- err = out
- elif not hasattr(err, 'write'):
- err = TextIO()
- self.err = err
- self.in_ = in_
- if now:
- self.startall()
- def startall(self):
- if self.out:
- sys.stdout = self.out
- if self.err:
- sys.stderr = self.err
- if self.in_:
- sys.stdin = self.in_ = DontReadFromInput()
- def done(self, save=True):
- """ return (outfile, errfile) and stop capturing. """
- outfile = errfile = None
- if self.out and not self.out.closed:
- sys.stdout = self._oldout
- outfile = self.out
- outfile.seek(0)
- if self.err and not self.err.closed:
- sys.stderr = self._olderr
- errfile = self.err
- errfile.seek(0)
- if self.in_:
- sys.stdin = self._oldin
- return outfile, errfile
- def resume(self):
- """ resume capturing with original temp files. """
- self.startall()
- def readouterr(self):
- """ return snapshot value of stdout/stderr capturings. """
- out = err = ""
- if self.out:
- out = self.out.getvalue()
- self.out.truncate(0)
- self.out.seek(0)
- if self.err:
- err = self.err.getvalue()
- self.err.truncate(0)
- self.err.seek(0)
- return out, err
- class DontReadFromInput:
- """Temporary stub class. Ideally when stdin is accessed, the
- capturing should be turned off, with possibly all data captured
- so far sent to the screen. This should be configurable, though,
- because in automated test runs it is better to crash than
- hang indefinitely.
- """
- def read(self, *args):
- raise IOError("reading from stdin while output is captured")
- readline = read
- readlines = read
- __iter__ = read
- def fileno(self):
- raise ValueError("redirected Stdin is pseudofile, has no fileno()")
- def isatty(self):
- return False
- def close(self):
- pass
- try:
- devnullpath = os.devnull
- except AttributeError:
- if os.name == 'nt':
- devnullpath = 'NUL'
- else:
- devnullpath = '/dev/null'
|