Fix ColorStream

This commit is contained in:
Todd Gamblin 2014-12-29 01:05:21 -08:00
parent 6ffcdc1166
commit a6e00f6086
2 changed files with 14 additions and 10 deletions

View File

@ -177,17 +177,20 @@ def cescape(string):
class ColorStream(object):
def __init__(self, stream, color=None):
self.__class__ = type(stream.__class__.__name__,
(self.__class__, stream.__class__), {})
self.__dict__ = stream.__dict__
self.color = color
self.stream = stream
self._stream = stream
self._color = color
def write(self, string, **kwargs):
if kwargs.get('raw', False):
super(ColorStream, self).write(string)
else:
cwrite(string, self.stream, self.color)
raw = kwargs.get('raw', False)
raw_write = getattr(self._stream, 'write')
color = self._color
if self._color is None:
if raw:
color=True
else:
color = self._stream.isatty()
raw_write(colorize(string, color=color))
def writelines(self, sequence, **kwargs):
raw = kwargs.get('raw', False)

View File

@ -1328,7 +1328,8 @@ def tree(self, **kwargs):
def graph(self, **kwargs):
N = kwargs.get('node', 'o') # Node character
out = kwargs.get('out', sys.stdout)
color = kwargs.get('color', True)
out = kwargs.get('out', ColorStream(sys.stdout, color=color))
indent = kwargs.get('indent', 0)
indent *= ' '