diff --git a/zipline/utils/log_utils.py b/zipline/utils/log_utils.py new file mode 100644 index 00000000..bf08cf41 --- /dev/null +++ b/zipline/utils/log_utils.py @@ -0,0 +1,161 @@ +import logbook +import zmq +import pytz +import datetime + +from logbook import NOTSET +from logbook.handlers import Handler + +from zipline.protocol import LOG_FRAME, LOG_FIELDS, \ + LOG_EXTRA_FIELDS + +from contextlib import contextmanager + +class redirecter(object): + def __init__(self, logger, name): + self.logger = logger + self.buffer = bytes() + self.name = name + + def write(self, line): + self.buffer += ''.join(['>>> ', line.strip('\n'), '\n']) + + def flush(self, final=False): + if not self.buffer: + return + out_form = """ [{pipe_name}] \n{buffer}""".format( + pipe_name = self.name, + buffer = self.buffer + ) + self.logger.error(out_form) + self.buffer = bytes() + +class log_redirecter(object): + def __init__(self, logger): + self.logger = logger + + def write(self, line): + #Absorb blank lines from print statements. + if line =='\n': + return + + else: + #TODO: add logic to guarantee we made this + self.logger.info(line.strip('\n')) + + def flush(self, final=False): + pass + +@contextmanager +def stdout_pipe(logger, pipe_name): + """ + Pipe stdout and stderr into a python logger interface + """ + import sys + orig_fds = sys.stdout, sys.stderr + + sys.stderr = redirecter(logger, pipe_name) + sys.stdout = redirecter(logger, pipe_name) + + yield + sys.stderr.flush() + sys.stdout.flush() + sys.stdout, sys.stderr = orig_fds + +@contextmanager +def stdout_only_pipe(logger, pipe_name): + """ + Pipes just stdout into a python logger interface + """ + import sys + orig_fd = sys.stdout + sys.stdout = log_redirecter(logger) + + yield + sys.stdout.flush() + sys.stdout = orig_fd + +class ZeroMQLogHandler(Handler): + """ + A handler that takes messages captured from the user algorithm stdout + and transforms them into LOG_FRAMES suitable for database storage. + Setup is similar to logbook.queues.ZeroMQHandler, except we connect + instead of binding and we extract record fields into a dict. + """ + + def __init__(self, uri=None, level=NOTSET, filter=None, bubble=False, + context=None, fds = LOG_FIELDS, extra_fds = LOG_EXTRA_FIELDS): + Handler.__init__(self, level, filter, bubble) + try: + import zmq + except ImportError: + raise RuntimeError('The pyzmq library is required for ' + 'the ZeroMQHandler.') + #: the zero mq context + self.context = context + #: the zero mq socket. + self.socket = self.context.socket(zmq.PUSH) + + self.uri = uri + if uri is not None: + self.socket.connect(uri) + + self.fds = fds + self.extra_fds = extra_fds + + def export_record(self, record): + """ + Extract relevant fields from a log record, fiddling with datetime + fields to make json happy. + """ + from zipline.utils.date_utils import EPOCH + + #Needed to extract record info from dictionary. + record.pull_information() + + #Logbook stores record times as datetime objects, which + #can't be serialized by JSON, so we need to convert to + #unix epoch representation. + + if record.time: + assert isinstance(record.time, datetime.datetime) + + time = record.time.replace(tzinfo = pytz.utc) + #logbook measures time in utc already, no need to convert. + record.time = EPOCH(time) + + #Do the same if algo_dt is a datetime object. + if record.extra.has_key('algo_dt'): + algo_dt = record.extra['algo_dt'] + + if isinstance(algo_dt, datetime.datetime): + algo_dt = EPOCH(algo_dt.replace(tzinfo = pytz.utc)) + record.extra['algo_dt'] = algo_dt + + data = {} + + #Extract all the fields we care about from LogRecord's internal + #dictionary. + + for field in iter(self.fds): + if record.__dict__.has_key(field): + data[field] = record.__dict__[field] + else: + data[field] = None + + for field in iter(self.extra_fds): + if record.extra.has_key(field): + data[field] = record.extra[field] + else: + data[field] = None + return data + + def emit(self, record): + """Extract relevant fields and send info as JSON over a zmq socket.""" + payload = self.export_record(record) + self.socket.send(LOG_FRAME(payload)) + + def close(self): + #self.socket.close() + pass +