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): class ColorStream(object):
def __init__(self, stream, color=None): def __init__(self, stream, color=None):
self.__class__ = type(stream.__class__.__name__, self._stream = stream
(self.__class__, stream.__class__), {}) self._color = color
self.__dict__ = stream.__dict__
self.color = color
self.stream = stream
def write(self, string, **kwargs): def write(self, string, **kwargs):
if kwargs.get('raw', False): raw = kwargs.get('raw', False)
super(ColorStream, self).write(string) raw_write = getattr(self._stream, 'write')
else:
cwrite(string, self.stream, self.color) 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): def writelines(self, sequence, **kwargs):
raw = kwargs.get('raw', False) raw = kwargs.get('raw', False)

View File

@ -1328,7 +1328,8 @@ def tree(self, **kwargs):
def graph(self, **kwargs): def graph(self, **kwargs):
N = kwargs.get('node', 'o') # Node character 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 = kwargs.get('indent', 0)
indent *= ' ' indent *= ' '