Merge pull request #5098 from minrk/parallel-debug

mostly debugging changes for IPython.parallel

IPython/parallel/controller/dictdb.py

@@ -1,317 +1,317 b''
 """A Task logger that presents our DB interface,
 but exists entirely in memory and implemented with dicts.
 
 Authors:
 
 * Min RK
 
 
 TaskRecords are dicts of the form::
 
     {
         'msg_id' : str(uuid),
         'client_uuid' : str(uuid),
         'engine_uuid' : str(uuid) or None,
         'header' : dict(header),
         'content': dict(content),
         'buffers': list(buffers),
-        'submitted': datetime,
+        'submitted': datetime or None,
         'started': datetime or None,
         'completed': datetime or None,
-        'resubmitted': datetime or None,
+        'received': datetime or None,
+        'resubmitted': str(uuid) or None,
         'result_header' : dict(header) or None,
         'result_content' : dict(content) or None,
         'result_buffers' : list(buffers) or None,
     }
 
 With this info, many of the special categories of tasks can be defined by query,
 e.g.:
 
 * pending: completed is None
 * client's outstanding: client_uuid = uuid && completed is None
 * MIA: arrived is None (and completed is None)
 
-EngineRecords are dicts of the form::
-
-    {
-        'eid' : int(id),
-        'uuid': str(uuid)
-    }
-
-This may be extended, but is currently.
-
-We support a subset of mongodb operators::
+DictDB supports a subset of mongodb operators::
 
     $lt,$gt,$lte,$gte,$ne,$in,$nin,$all,$mod,$exists
 """
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2010-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 from copy import deepcopy as copy
 from datetime import datetime
 
 from IPython.config.configurable import LoggingConfigurable
 
 from IPython.utils.py3compat import iteritems, itervalues
 from IPython.utils.traitlets import Dict, Unicode, Integer, Float
 
 filters = {
  '$lt' : lambda a,b: a < b,
  '$gt' : lambda a,b: b > a,
  '$eq' : lambda a,b: a == b,
  '$ne' : lambda a,b: a != b,
  '$lte': lambda a,b: a <= b,
  '$gte': lambda a,b: a >= b,
  '$in' : lambda a,b: a in b,
  '$nin': lambda a,b: a not in b,
  '$all': lambda a,b: all([ a in bb for bb in b ]),
  '$mod': lambda a,b: a%b[0] == b[1],
  '$exists' : lambda a,b: (b and a is not None) or (a is None and not b)
 }
 
 
 class CompositeFilter(object):
     """Composite filter for matching multiple properties."""
 
     def __init__(self, dikt):
         self.tests = []
         self.values = []
         for key, value in iteritems(dikt):
             self.tests.append(filters[key])
             self.values.append(value)
 
     def __call__(self, value):
         for test,check in zip(self.tests, self.values):
             if not test(value, check):
                 return False
         return True
 
 class BaseDB(LoggingConfigurable):
     """Empty Parent class so traitlets work on DB."""
     # base configurable traits:
     session = Unicode("")
 
 class DictDB(BaseDB):
     """Basic in-memory dict-based object for saving Task Records.
 
     This is the first object to present the DB interface
     for logging tasks out of memory.
 
     The interface is based on MongoDB, so adding a MongoDB
     backend should be straightforward.
     """
 
     _records = Dict()
     _culled_ids = set() # set of ids which have been culled
     _buffer_bytes = Integer(0) # running total of the bytes in the DB
     
     size_limit = Integer(1024**3, config=True,
         help="""The maximum total size (in bytes) of the buffers stored in the db
         
         When the db exceeds this size, the oldest records will be culled until
         the total size is under size_limit * (1-cull_fraction).
         default: 1 GB
         """
     )
     record_limit = Integer(1024, config=True,
         help="""The maximum number of records in the db
         
         When the history exceeds this size, the first record_limit * cull_fraction
         records will be culled.
         """
     )
     cull_fraction = Float(0.1, config=True,
         help="""The fraction by which the db should culled when one of the limits is exceeded
         
         In general, the db size will spend most of its time with a size in the range:
         
         [limit * (1-cull_fraction), limit]
         
         for each of size_limit and record_limit.
         """
     )
 
     def _match_one(self, rec, tests):
         """Check if a specific record matches tests."""
         for key,test in iteritems(tests):
             if not test(rec.get(key, None)):
                 return False
         return True
 
     def _match(self, check):
         """Find all the matches for a check dict."""
         matches = []
         tests = {}
         for k,v in iteritems(check):
             if isinstance(v, dict):
                 tests[k] = CompositeFilter(v)
             else:
                 tests[k] = lambda o: o==v
 
         for rec in itervalues(self._records):
             if self._match_one(rec, tests):
                 matches.append(copy(rec))
         return matches
 
     def _extract_subdict(self, rec, keys):
         """extract subdict of keys"""
         d = {}
         d['msg_id'] = rec['msg_id']
         for key in keys:
             d[key] = rec[key]
         return copy(d)
     
     # methods for monitoring size / culling history
     
     def _add_bytes(self, rec):
         for key in ('buffers', 'result_buffers'):
             for buf in rec.get(key) or []:
                 self._buffer_bytes += len(buf)
         
         self._maybe_cull()
     
     def _drop_bytes(self, rec):
         for key in ('buffers', 'result_buffers'):
             for buf in rec.get(key) or []:
                 self._buffer_bytes -= len(buf)
     
     def _cull_oldest(self, n=1):
         """cull the oldest N records"""
         for msg_id in self.get_history()[:n]:
             self.log.debug("Culling record: %r", msg_id)
             self._culled_ids.add(msg_id)
             self.drop_record(msg_id)
     
     def _maybe_cull(self):
         # cull by count:
         if len(self._records) > self.record_limit:
             to_cull = int(self.cull_fraction * self.record_limit)
             self.log.info("%i records exceeds limit of %i, culling oldest %i",
                 len(self._records), self.record_limit, to_cull
             )
             self._cull_oldest(to_cull)
         
         # cull by size:
         if self._buffer_bytes > self.size_limit:
             limit = self.size_limit * (1 - self.cull_fraction)
             
             before = self._buffer_bytes
             before_count = len(self._records)
             culled = 0
             while self._buffer_bytes > limit:
                 self._cull_oldest(1)
                 culled += 1
         
             self.log.info("%i records with total buffer size %i exceeds limit: %i. Culled oldest %i records.",
                 before_count, before, self.size_limit, culled
             )
     
+    def _check_dates(self, rec):
+        for key in ('submitted', 'started', 'completed'):
+            value = rec.get(key, None)
+            if value is not None and not isinstance(value, datetime):
+                raise ValueError("%s must be None or datetime, not %r" % (key, value))
+    
     # public API methods:
 
     def add_record(self, msg_id, rec):
         """Add a new Task Record, by msg_id."""
         if msg_id in self._records:
             raise KeyError("Already have msg_id %r"%(msg_id))
+        self._check_dates(rec)
         self._records[msg_id] = rec
         self._add_bytes(rec)
         self._maybe_cull()
 
     def get_record(self, msg_id):
         """Get a specific Task Record, by msg_id."""
         if msg_id in self._culled_ids:
             raise KeyError("Record %r has been culled for size" % msg_id)
         if not msg_id in self._records:
             raise KeyError("No such msg_id %r"%(msg_id))
         return copy(self._records[msg_id])
 
     def update_record(self, msg_id, rec):
         """Update the data in an existing record."""
         if msg_id in self._culled_ids:
             raise KeyError("Record %r has been culled for size" % msg_id)
+        self._check_dates(rec)
         _rec = self._records[msg_id]
         self._drop_bytes(_rec)
         _rec.update(rec)
         self._add_bytes(_rec)
 
     def drop_matching_records(self, check):
         """Remove a record from the DB."""
         matches = self._match(check)
         for rec in matches:
             self._drop_bytes(rec)
             del self._records[rec['msg_id']]
 
     def drop_record(self, msg_id):
         """Remove a record from the DB."""
         rec = self._records[msg_id]
         self._drop_bytes(rec)
         del self._records[msg_id]
 
     def find_records(self, check, keys=None):
         """Find records matching a query dict, optionally extracting subset of keys.
 
         Returns dict keyed by msg_id of matching records.
 
         Parameters
         ----------
 
         check: dict
             mongodb-style query argument
         keys: list of strs [optional]
             if specified, the subset of keys to extract.  msg_id will *always* be
             included.
         """
         matches = self._match(check)
         if keys:
             return [ self._extract_subdict(rec, keys) for rec in matches ]
         else:
             return matches
 
     def get_history(self):
         """get all msg_ids, ordered by time submitted."""
         msg_ids = self._records.keys()
         # Remove any that do not have a submitted timestamp.
         # This is extremely unlikely to happen,
         # but it seems to come up in some tests on VMs.
         msg_ids = [ m for m in msg_ids if self._records[m]['submitted'] is not None ]
         return sorted(msg_ids, key=lambda m: self._records[m]['submitted'])
 
 
 NODATA = KeyError("NoDB backend doesn't store any data. "
 "Start the Controller with a DB backend to enable resubmission / result persistence."
 )
 
 
 class NoDB(BaseDB):
     """A blackhole db backend that actually stores no information.
     
     Provides the full DB interface, but raises KeyErrors on any
     method that tries to access the records.  This can be used to
     minimize the memory footprint of the Hub when its record-keeping
     functionality is not required.
     """
     
     def add_record(self, msg_id, record):
         pass
     
     def get_record(self, msg_id):
         raise NODATA
     
     def update_record(self, msg_id, record):
         pass
     
     def drop_matching_records(self, check):
         pass
     
     def drop_record(self, msg_id):
         pass
     
     def find_records(self, check, keys=None):
         raise NODATA
     
     def get_history(self):
         raise NODATA
 

IPython/parallel/controller/hub.py

@@ -1,1426 +1,1436 b''
 """The IPython Controller Hub with 0MQ
 This is the master object that handles connections from engines and clients,
 and monitors traffic through the various queues.
 
 Authors:
 
 * Min RK
 """
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2010-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 from __future__ import print_function
 
 import json
 import os
 import sys
 import time
 from datetime import datetime
 
 import zmq
 from zmq.eventloop import ioloop
 from zmq.eventloop.zmqstream import ZMQStream
 
 # internal:
 from IPython.utils.importstring import import_item
 from IPython.utils.jsonutil import extract_dates
 from IPython.utils.localinterfaces import localhost
 from IPython.utils.py3compat import cast_bytes, unicode_type, iteritems
 from IPython.utils.traitlets import (
         HasTraits, Instance, Integer, Unicode, Dict, Set, Tuple, CBytes, DottedObjectName
         )
 
 from IPython.parallel import error, util
 from IPython.parallel.factory import RegistrationFactory
 
 from IPython.kernel.zmq.session import SessionFactory
 
 from .heartmonitor import HeartMonitor
 
 #-----------------------------------------------------------------------------
 # Code
 #-----------------------------------------------------------------------------
 
 def _passer(*args, **kwargs):
     return
 
 def _printer(*args, **kwargs):
     print (args)
     print (kwargs)
 
 def empty_record():
     """Return an empty dict with all record keys."""
     return {
         'msg_id' : None,
         'header' : None,
         'metadata' : None,
         'content': None,
         'buffers': None,
         'submitted': None,
         'client_uuid' : None,
         'engine_uuid' : None,
         'started': None,
         'completed': None,
         'resubmitted': None,
         'received': None,
         'result_header' : None,
         'result_metadata' : None,
         'result_content' : None,
         'result_buffers' : None,
         'queue' : None,
         'pyin' : None,
         'pyout': None,
         'pyerr': None,
         'stdout': '',
         'stderr': '',
     }
 
 def init_record(msg):
     """Initialize a TaskRecord based on a request."""
     header = msg['header']
     return {
         'msg_id' : header['msg_id'],
         'header' : header,
         'content': msg['content'],
         'metadata': msg['metadata'],
         'buffers': msg['buffers'],
         'submitted': header['date'],
         'client_uuid' : None,
         'engine_uuid' : None,
         'started': None,
         'completed': None,
         'resubmitted': None,
         'received': None,
         'result_header' : None,
         'result_metadata': None,
         'result_content' : None,
         'result_buffers' : None,
         'queue' : None,
         'pyin' : None,
         'pyout': None,
         'pyerr': None,
         'stdout': '',
         'stderr': '',
     }
 
 
 class EngineConnector(HasTraits):
     """A simple object for accessing the various zmq connections of an object.
     Attributes are:
     id (int): engine ID
     uuid (unicode): engine UUID
     pending: set of msg_ids
     stallback: DelayedCallback for stalled registration
     """
     
     id = Integer(0)
     uuid = Unicode()
     pending = Set()
     stallback = Instance(ioloop.DelayedCallback)
 
 
 _db_shortcuts = {
     'sqlitedb' : 'IPython.parallel.controller.sqlitedb.SQLiteDB',
     'mongodb'  : 'IPython.parallel.controller.mongodb.MongoDB',
     'dictdb'   : 'IPython.parallel.controller.dictdb.DictDB',
     'nodb'     : 'IPython.parallel.controller.dictdb.NoDB',
 }
 
 class HubFactory(RegistrationFactory):
     """The Configurable for setting up a Hub."""
 
     # port-pairs for monitoredqueues:
     hb = Tuple(Integer,Integer,config=True,
         help="""PUB/ROUTER Port pair for Engine heartbeats""")
     def _hb_default(self):
         return tuple(util.select_random_ports(2))
 
     mux = Tuple(Integer,Integer,config=True,
         help="""Client/Engine Port pair for MUX queue""")
 
     def _mux_default(self):
         return tuple(util.select_random_ports(2))
 
     task = Tuple(Integer,Integer,config=True,
         help="""Client/Engine Port pair for Task queue""")
     def _task_default(self):
         return tuple(util.select_random_ports(2))
 
     control = Tuple(Integer,Integer,config=True,
         help="""Client/Engine Port pair for Control queue""")
 
     def _control_default(self):
         return tuple(util.select_random_ports(2))
 
     iopub = Tuple(Integer,Integer,config=True,
         help="""Client/Engine Port pair for IOPub relay""")
 
     def _iopub_default(self):
         return tuple(util.select_random_ports(2))
 
     # single ports:
     mon_port = Integer(config=True,
         help="""Monitor (SUB) port for queue traffic""")
 
     def _mon_port_default(self):
         return util.select_random_ports(1)[0]
 
     notifier_port = Integer(config=True,
         help="""PUB port for sending engine status notifications""")
 
     def _notifier_port_default(self):
         return util.select_random_ports(1)[0]
 
     engine_ip = Unicode(config=True,
         help="IP on which to listen for engine connections. [default: loopback]")
     def _engine_ip_default(self):
         return localhost()
     engine_transport = Unicode('tcp', config=True,
         help="0MQ transport for engine connections. [default: tcp]")
 
     client_ip = Unicode(config=True,
         help="IP on which to listen for client connections. [default: loopback]")
     client_transport = Unicode('tcp', config=True,
         help="0MQ transport for client connections. [default : tcp]")
 
     monitor_ip = Unicode(config=True,
         help="IP on which to listen for monitor messages. [default: loopback]")
     monitor_transport = Unicode('tcp', config=True,
         help="0MQ transport for monitor messages. [default : tcp]")
 
     _client_ip_default = _monitor_ip_default = _engine_ip_default
 
 
     monitor_url = Unicode('')
 
     db_class = DottedObjectName('NoDB',
         config=True, help="""The class to use for the DB backend
         
         Options include:
         
         SQLiteDB: SQLite
         MongoDB : use MongoDB
         DictDB  : in-memory storage (fastest, but be mindful of memory growth of the Hub)
         NoDB    : disable database altogether (default)
         
         """)
 
     # not configurable
     db = Instance('IPython.parallel.controller.dictdb.BaseDB')
     heartmonitor = Instance('IPython.parallel.controller.heartmonitor.HeartMonitor')
 
     def _ip_changed(self, name, old, new):
         self.engine_ip = new
         self.client_ip = new
         self.monitor_ip = new
         self._update_monitor_url()
 
     def _update_monitor_url(self):
         self.monitor_url = "%s://%s:%i" % (self.monitor_transport, self.monitor_ip, self.mon_port)
 
     def _transport_changed(self, name, old, new):
         self.engine_transport = new
         self.client_transport = new
         self.monitor_transport = new
         self._update_monitor_url()
 
     def __init__(self, **kwargs):
         super(HubFactory, self).__init__(**kwargs)
         self._update_monitor_url()
 
 
     def construct(self):
         self.init_hub()
 
     def start(self):
         self.heartmonitor.start()
         self.log.info("Heartmonitor started")
 
     def client_url(self, channel):
         """return full zmq url for a named client channel"""
         return "%s://%s:%i" % (self.client_transport, self.client_ip, self.client_info[channel])
     
     def engine_url(self, channel):
         """return full zmq url for a named engine channel"""
         return "%s://%s:%i" % (self.engine_transport, self.engine_ip, self.engine_info[channel])
     
     def init_hub(self):
         """construct Hub object"""
 
         ctx = self.context
         loop = self.loop
         if 'TaskScheduler.scheme_name' in self.config:
             scheme = self.config.TaskScheduler.scheme_name
         else:
             from .scheduler import TaskScheduler
             scheme = TaskScheduler.scheme_name.get_default_value()
         
         # build connection dicts
         engine = self.engine_info = {
             'interface'     : "%s://%s" % (self.engine_transport, self.engine_ip),
             'registration'  : self.regport,
             'control'       : self.control[1],
             'mux'           : self.mux[1],
             'hb_ping'       : self.hb[0],
             'hb_pong'       : self.hb[1],
             'task'          : self.task[1],
             'iopub'         : self.iopub[1],
             }
 
         client = self.client_info = {
             'interface'     : "%s://%s" % (self.client_transport, self.client_ip),
             'registration'  : self.regport,
             'control'       : self.control[0],
             'mux'           : self.mux[0],
             'task'          : self.task[0],
             'task_scheme'   : scheme,
             'iopub'         : self.iopub[0],
             'notification'  : self.notifier_port,
             }
         
         self.log.debug("Hub engine addrs: %s", self.engine_info)
         self.log.debug("Hub client addrs: %s", self.client_info)
         
         # Registrar socket
         q = ZMQStream(ctx.socket(zmq.ROUTER), loop)
         util.set_hwm(q, 0)
         q.bind(self.client_url('registration'))
         self.log.info("Hub listening on %s for registration.", self.client_url('registration'))
         if self.client_ip != self.engine_ip:
             q.bind(self.engine_url('registration'))
             self.log.info("Hub listening on %s for registration.", self.engine_url('registration'))
 
         ### Engine connections ###
 
         # heartbeat
         hpub = ctx.socket(zmq.PUB)
         hpub.bind(self.engine_url('hb_ping'))
         hrep = ctx.socket(zmq.ROUTER)
         util.set_hwm(hrep, 0)
         hrep.bind(self.engine_url('hb_pong'))
         self.heartmonitor = HeartMonitor(loop=loop, parent=self, log=self.log,
                                 pingstream=ZMQStream(hpub,loop),
                                 pongstream=ZMQStream(hrep,loop)
                             )
 
         ### Client connections ###
         
         # Notifier socket
         n = ZMQStream(ctx.socket(zmq.PUB), loop)
         n.bind(self.client_url('notification'))
 
         ### build and launch the queues ###
 
         # monitor socket
         sub = ctx.socket(zmq.SUB)
         sub.setsockopt(zmq.SUBSCRIBE, b"")
         sub.bind(self.monitor_url)
         sub.bind('inproc://monitor')
         sub = ZMQStream(sub, loop)
 
         # connect the db
         db_class = _db_shortcuts.get(self.db_class.lower(), self.db_class)
         self.log.info('Hub using DB backend: %r', (db_class.split('.')[-1]))
         self.db = import_item(str(db_class))(session=self.session.session,
                                             parent=self, log=self.log)
         time.sleep(.25)
 
         # resubmit stream
         r = ZMQStream(ctx.socket(zmq.DEALER), loop)
         url = util.disambiguate_url(self.client_url('task'))
         r.connect(url)
 
         self.hub = Hub(loop=loop, session=self.session, monitor=sub, heartmonitor=self.heartmonitor,
                 query=q, notifier=n, resubmit=r, db=self.db,
                 engine_info=self.engine_info, client_info=self.client_info,
                 log=self.log)
 
 
 class Hub(SessionFactory):
     """The IPython Controller Hub with 0MQ connections
 
     Parameters
     ==========
     loop: zmq IOLoop instance
     session: Session object
     <removed> context: zmq context for creating new connections (?)
     queue: ZMQStream for monitoring the command queue (SUB)
     query: ZMQStream for engine registration and client queries requests (ROUTER)
     heartbeat: HeartMonitor object checking the pulse of the engines
     notifier: ZMQStream for broadcasting engine registration changes (PUB)
     db: connection to db for out of memory logging of commands
                 NotImplemented
     engine_info: dict of zmq connection information for engines to connect
                 to the queues.
     client_info: dict of zmq connection information for engines to connect
                 to the queues.
     """
     
     engine_state_file = Unicode()
     
     # internal data structures:
     ids=Set() # engine IDs
     keytable=Dict()
     by_ident=Dict()
     engines=Dict()
     clients=Dict()
     hearts=Dict()
     pending=Set()
     queues=Dict()  # pending msg_ids keyed by engine_id
     tasks=Dict() # pending msg_ids submitted as tasks, keyed by client_id
     completed=Dict() # completed msg_ids keyed by engine_id
     all_completed=Set() # completed msg_ids keyed by engine_id
     dead_engines=Set() # completed msg_ids keyed by engine_id
     unassigned=Set() # set of task msg_ds not yet assigned a destination
     incoming_registrations=Dict()
     registration_timeout=Integer()
     _idcounter=Integer(0)
 
     # objects from constructor:
     query=Instance(ZMQStream)
     monitor=Instance(ZMQStream)
     notifier=Instance(ZMQStream)
     resubmit=Instance(ZMQStream)
     heartmonitor=Instance(HeartMonitor)
     db=Instance(object)
     client_info=Dict()
     engine_info=Dict()
 
 
     def __init__(self, **kwargs):
         """
         # universal:
         loop: IOLoop for creating future connections
         session: streamsession for sending serialized data
         # engine:
         queue: ZMQStream for monitoring queue messages
         query: ZMQStream for engine+client registration and client requests
         heartbeat: HeartMonitor object for tracking engines
         # extra:
         db: ZMQStream for db connection (NotImplemented)
         engine_info: zmq address/protocol dict for engine connections
         client_info: zmq address/protocol dict for client connections
         """
 
         super(Hub, self).__init__(**kwargs)
         self.registration_timeout = max(10000, 5*self.heartmonitor.period)
 
         # register our callbacks
         self.query.on_recv(self.dispatch_query)
         self.monitor.on_recv(self.dispatch_monitor_traffic)
 
         self.heartmonitor.add_heart_failure_handler(self.handle_heart_failure)
         self.heartmonitor.add_new_heart_handler(self.handle_new_heart)
 
         self.monitor_handlers = {b'in' : self.save_queue_request,
                                 b'out': self.save_queue_result,
                                 b'intask': self.save_task_request,
                                 b'outtask': self.save_task_result,
                                 b'tracktask': self.save_task_destination,
                                 b'incontrol': _passer,
                                 b'outcontrol': _passer,
                                 b'iopub': self.save_iopub_message,
         }
 
         self.query_handlers = {'queue_request': self.queue_status,
                                 'result_request': self.get_results,
                                 'history_request': self.get_history,
                                 'db_request': self.db_query,
                                 'purge_request': self.purge_results,
                                 'load_request': self.check_load,
                                 'resubmit_request': self.resubmit_task,
                                 'shutdown_request': self.shutdown_request,
                                 'registration_request' : self.register_engine,
                                 'unregistration_request' : self.unregister_engine,
                                 'connection_request': self.connection_request,
         }
 
         # ignore resubmit replies
         self.resubmit.on_recv(lambda msg: None, copy=False)
 
         self.log.info("hub::created hub")
     
     @property
     def _next_id(self):
         """gemerate a new ID.
 
         No longer reuse old ids, just count from 0."""
         newid = self._idcounter
         self._idcounter += 1
         return newid
         # newid = 0
         # incoming = [id[0] for id in itervalues(self.incoming_registrations)]
         # # print newid, self.ids, self.incoming_registrations
         # while newid in self.ids or newid in incoming:
         #     newid += 1
         # return newid
     
     #-----------------------------------------------------------------------------
     # message validation
     #-----------------------------------------------------------------------------
 
     def _validate_targets(self, targets):
         """turn any valid targets argument into a list of integer ids"""
         if targets is None:
             # default to all
             return self.ids
 
         if isinstance(targets, (int,str,unicode_type)):
             # only one target specified
             targets = [targets]
         _targets = []
         for t in targets:
             # map raw identities to ids
             if isinstance(t, (str,unicode_type)):
                 t = self.by_ident.get(cast_bytes(t), t)
             _targets.append(t)
         targets = _targets
         bad_targets = [ t for t in targets if t not in self.ids ]
         if bad_targets:
             raise IndexError("No Such Engine: %r" % bad_targets)
         if not targets:
             raise IndexError("No Engines Registered")
         return targets
 
     #-----------------------------------------------------------------------------
     # dispatch methods (1 per stream)
     #-----------------------------------------------------------------------------
 
 
     @util.log_errors
     def dispatch_monitor_traffic(self, msg):
         """all ME and Task queue messages come through here, as well as
         IOPub traffic."""
         self.log.debug("monitor traffic: %r", msg[0])
         switch = msg[0]
         try:
             idents, msg = self.session.feed_identities(msg[1:])
         except ValueError:
             idents=[]
         if not idents:
             self.log.error("Monitor message without topic: %r", msg)
             return
         handler = self.monitor_handlers.get(switch, None)
         if handler is not None:
             handler(idents, msg)
         else:
             self.log.error("Unrecognized monitor topic: %r", switch)
 
 
     @util.log_errors
     def dispatch_query(self, msg):
         """Route registration requests and queries from clients."""
         try:
             idents, msg = self.session.feed_identities(msg)
         except ValueError:
             idents = []
         if not idents:
             self.log.error("Bad Query Message: %r", msg)
             return
         client_id = idents[0]
         try:
             msg = self.session.unserialize(msg, content=True)
         except Exception:
             content = error.wrap_exception()
             self.log.error("Bad Query Message: %r", msg, exc_info=True)
             self.session.send(self.query, "hub_error", ident=client_id,
                     content=content)
             return
         # print client_id, header, parent, content
         #switch on message type:
         msg_type = msg['header']['msg_type']
         self.log.info("client::client %r requested %r", client_id, msg_type)
         handler = self.query_handlers.get(msg_type, None)
         try:
             assert handler is not None, "Bad Message Type: %r" % msg_type
         except:
             content = error.wrap_exception()
             self.log.error("Bad Message Type: %r", msg_type, exc_info=True)
             self.session.send(self.query, "hub_error", ident=client_id,
                     content=content)
             return
 
         else:
             handler(idents, msg)
 
     def dispatch_db(self, msg):
         """"""
         raise NotImplementedError
 
     #---------------------------------------------------------------------------
     # handler methods (1 per event)
     #---------------------------------------------------------------------------
 
     #----------------------- Heartbeat --------------------------------------
 
     def handle_new_heart(self, heart):
         """handler to attach to heartbeater.
         Called when a new heart starts to beat.
         Triggers completion of registration."""
         self.log.debug("heartbeat::handle_new_heart(%r)", heart)
         if heart not in self.incoming_registrations:
             self.log.info("heartbeat::ignoring new heart: %r", heart)
         else:
             self.finish_registration(heart)
 
 
     def handle_heart_failure(self, heart):
         """handler to attach to heartbeater.
         called when a previously registered heart fails to respond to beat request.
         triggers unregistration"""
         self.log.debug("heartbeat::handle_heart_failure(%r)", heart)
         eid = self.hearts.get(heart, None)
         uuid = self.engines[eid].uuid
         if eid is None or self.keytable[eid] in self.dead_engines:
             self.log.info("heartbeat::ignoring heart failure %r (not an engine or already dead)", heart)
         else:
             self.unregister_engine(heart, dict(content=dict(id=eid, queue=uuid)))
 
     #----------------------- MUX Queue Traffic ------------------------------
 
     def save_queue_request(self, idents, msg):
         if len(idents) < 2:
             self.log.error("invalid identity prefix: %r", idents)
             return
         queue_id, client_id = idents[:2]
         try:
             msg = self.session.unserialize(msg)
         except Exception:
             self.log.error("queue::client %r sent invalid message to %r: %r", client_id, queue_id, msg, exc_info=True)
             return
 
         eid = self.by_ident.get(queue_id, None)
         if eid is None:
             self.log.error("queue::target %r not registered", queue_id)
             self.log.debug("queue::    valid are: %r", self.by_ident.keys())
             return
         record = init_record(msg)
         msg_id = record['msg_id']
         self.log.info("queue::client %r submitted request %r to %s", client_id, msg_id, eid)
         # Unicode in records
         record['engine_uuid'] = queue_id.decode('ascii')
         record['client_uuid'] = msg['header']['session']
         record['queue'] = 'mux'
 
         try:
             # it's posible iopub arrived first:
             existing = self.db.get_record(msg_id)
             for key,evalue in iteritems(existing):
                 rvalue = record.get(key, None)
                 if evalue and rvalue and evalue != rvalue:
                     self.log.warn("conflicting initial state for record: %r:%r <%r> %r", msg_id, rvalue, key, evalue)
                 elif evalue and not rvalue:
                     record[key] = evalue
             try:
                 self.db.update_record(msg_id, record)
             except Exception:
                 self.log.error("DB Error updating record %r", msg_id, exc_info=True)
         except KeyError:
             try:
                 self.db.add_record(msg_id, record)
             except Exception:
                 self.log.error("DB Error adding record %r", msg_id, exc_info=True)
 
 
         self.pending.add(msg_id)
         self.queues[eid].append(msg_id)
 
     def save_queue_result(self, idents, msg):
         if len(idents) < 2:
             self.log.error("invalid identity prefix: %r", idents)
             return
 
         client_id, queue_id = idents[:2]
         try:
             msg = self.session.unserialize(msg)
         except Exception:
             self.log.error("queue::engine %r sent invalid message to %r: %r",
                     queue_id, client_id, msg, exc_info=True)
             return
 
         eid = self.by_ident.get(queue_id, None)
         if eid is None:
             self.log.error("queue::unknown engine %r is sending a reply: ", queue_id)
             return
 
         parent = msg['parent_header']
         if not parent:
             return
         msg_id = parent['msg_id']
         if msg_id in self.pending:
             self.pending.remove(msg_id)
             self.all_completed.add(msg_id)
             self.queues[eid].remove(msg_id)
             self.completed[eid].append(msg_id)
             self.log.info("queue::request %r completed on %s", msg_id, eid)
         elif msg_id not in self.all_completed:
             # it could be a result from a dead engine that died before delivering the
             # result
             self.log.warn("queue:: unknown msg finished %r", msg_id)
             return
         # update record anyway, because the unregistration could have been premature
         rheader = msg['header']
         md = msg['metadata']
         completed = rheader['date']
-        started = md.get('started', None)
+        started = extract_dates(md.get('started', None))
         result = {
             'result_header' : rheader,
             'result_metadata': md,
             'result_content': msg['content'],
             'received': datetime.now(),
             'started' : started,
             'completed' : completed
         }
 
         result['result_buffers'] = msg['buffers']
         try:
             self.db.update_record(msg_id, result)
         except Exception:
             self.log.error("DB Error updating record %r", msg_id, exc_info=True)
 
 
     #--------------------- Task Queue Traffic ------------------------------
 
     def save_task_request(self, idents, msg):
         """Save the submission of a task."""
         client_id = idents[0]
 
         try:
             msg = self.session.unserialize(msg)
         except Exception:
             self.log.error("task::client %r sent invalid task message: %r",
                     client_id, msg, exc_info=True)
             return
         record = init_record(msg)
 
         record['client_uuid'] = msg['header']['session']
         record['queue'] = 'task'
         header = msg['header']
         msg_id = header['msg_id']
         self.pending.add(msg_id)
         self.unassigned.add(msg_id)
         try:
             # it's posible iopub arrived first:
             existing = self.db.get_record(msg_id)
             if existing['resubmitted']:
                 for key in ('submitted', 'client_uuid', 'buffers'):
                     # don't clobber these keys on resubmit
                     # submitted and client_uuid should be different
                     # and buffers might be big, and shouldn't have changed
                     record.pop(key)
                     # still check content,header which should not change
                     # but are not expensive to compare as buffers
 
             for key,evalue in iteritems(existing):
                 if key.endswith('buffers'):
                     # don't compare buffers
                     continue
                 rvalue = record.get(key, None)
                 if evalue and rvalue and evalue != rvalue:
                     self.log.warn("conflicting initial state for record: %r:%r <%r> %r", msg_id, rvalue, key, evalue)
                 elif evalue and not rvalue:
                     record[key] = evalue
             try:
                 self.db.update_record(msg_id, record)
             except Exception:
                 self.log.error("DB Error updating record %r", msg_id, exc_info=True)
         except KeyError:
             try:
                 self.db.add_record(msg_id, record)
             except Exception:
                 self.log.error("DB Error adding record %r", msg_id, exc_info=True)
         except Exception:
             self.log.error("DB Error saving task request %r", msg_id, exc_info=True)
 
     def save_task_result(self, idents, msg):
         """save the result of a completed task."""
         client_id = idents[0]
         try:
             msg = self.session.unserialize(msg)
         except Exception:
             self.log.error("task::invalid task result message send to %r: %r",
                     client_id, msg, exc_info=True)
             return
 
         parent = msg['parent_header']
         if not parent:
             # print msg
             self.log.warn("Task %r had no parent!", msg)
             return
         msg_id = parent['msg_id']
         if msg_id in self.unassigned:
             self.unassigned.remove(msg_id)
 
         header = msg['header']
         md = msg['metadata']
         engine_uuid = md.get('engine', u'')
         eid = self.by_ident.get(cast_bytes(engine_uuid), None)
         
         status = md.get('status', None)
 
         if msg_id in self.pending:
             self.log.info("task::task %r finished on %s", msg_id, eid)
             self.pending.remove(msg_id)
             self.all_completed.add(msg_id)
             if eid is not None:
                 if status != 'aborted':
                     self.completed[eid].append(msg_id)
                 if msg_id in self.tasks[eid]:
                     self.tasks[eid].remove(msg_id)
             completed = header['date']
-            started = md.get('started', None)
+            started = extract_dates(md.get('started', None))
             result = {
                 'result_header' : header,
                 'result_metadata': msg['metadata'],
                 'result_content': msg['content'],
                 'started' : started,
                 'completed' : completed,
                 'received' : datetime.now(),
                 'engine_uuid': engine_uuid,
             }
 
             result['result_buffers'] = msg['buffers']
             try:
                 self.db.update_record(msg_id, result)
             except Exception:
                 self.log.error("DB Error saving task request %r", msg_id, exc_info=True)
 
         else:
             self.log.debug("task::unknown task %r finished", msg_id)
 
     def save_task_destination(self, idents, msg):
         try:
             msg = self.session.unserialize(msg, content=True)
         except Exception:
             self.log.error("task::invalid task tracking message", exc_info=True)
             return
         content = msg['content']
         # print (content)
         msg_id = content['msg_id']
         engine_uuid = content['engine_id']
         eid = self.by_ident[cast_bytes(engine_uuid)]
 
         self.log.info("task::task %r arrived on %r", msg_id, eid)
         if msg_id in self.unassigned:
             self.unassigned.remove(msg_id)
         # else:
         #     self.log.debug("task::task %r not listed as MIA?!"%(msg_id))
 
         self.tasks[eid].append(msg_id)
         # self.pending[msg_id][1].update(received=datetime.now(),engine=(eid,engine_uuid))
         try:
             self.db.update_record(msg_id, dict(engine_uuid=engine_uuid))
         except Exception:
             self.log.error("DB Error saving task destination %r", msg_id, exc_info=True)
 
 
     def mia_task_request(self, idents, msg):
         raise NotImplementedError
         client_id = idents[0]
         # content = dict(mia=self.mia,status='ok')
         # self.session.send('mia_reply', content=content, idents=client_id)
 
 
     #--------------------- IOPub Traffic ------------------------------
 
     def save_iopub_message(self, topics, msg):
         """save an iopub message into the db"""
         # print (topics)
         try:
             msg = self.session.unserialize(msg, content=True)
         except Exception:
             self.log.error("iopub::invalid IOPub message", exc_info=True)
             return
 
         parent = msg['parent_header']
         if not parent:
             self.log.warn("iopub::IOPub message lacks parent: %r", msg)
             return
         msg_id = parent['msg_id']
         msg_type = msg['header']['msg_type']
         content = msg['content']
 
         # ensure msg_id is in db
         try:
             rec = self.db.get_record(msg_id)
         except KeyError:
             rec = empty_record()
             rec['msg_id'] = msg_id
             self.db.add_record(msg_id, rec)
         # stream
         d = {}
         if msg_type == 'stream':
             name = content['name']
             s = rec[name] or ''
             d[name] = s + content['data']
 
         elif msg_type == 'pyerr':
             d['pyerr'] = content
         elif msg_type == 'pyin':
             d['pyin'] = content['code']
         elif msg_type in ('display_data', 'pyout'):
             d[msg_type] = content
         elif msg_type == 'status':
             pass
         elif msg_type == 'data_pub':
             self.log.info("ignored data_pub message for %s" % msg_id)
         else:
             self.log.warn("unhandled iopub msg_type: %r", msg_type)
 
         if not d:
             return
 
         try:
             self.db.update_record(msg_id, d)
         except Exception:
             self.log.error("DB Error saving iopub message %r", msg_id, exc_info=True)
 
 
 
     #-------------------------------------------------------------------------
     # Registration requests
     #-------------------------------------------------------------------------
 
     def connection_request(self, client_id, msg):
         """Reply with connection addresses for clients."""
         self.log.info("client::client %r connected", client_id)
         content = dict(status='ok')
         jsonable = {}
         for k,v in iteritems(self.keytable):
             if v not in self.dead_engines:
                 jsonable[str(k)] = v
         content['engines'] = jsonable
         self.session.send(self.query, 'connection_reply', content, parent=msg, ident=client_id)
 
     def register_engine(self, reg, msg):
         """Register a new engine."""
         content = msg['content']
         try:
             uuid = content['uuid']
         except KeyError:
             self.log.error("registration::queue not specified", exc_info=True)
             return
 
         eid = self._next_id
 
         self.log.debug("registration::register_engine(%i, %r)", eid, uuid)
 
         content = dict(id=eid,status='ok',hb_period=self.heartmonitor.period)
         # check if requesting available IDs:
         if cast_bytes(uuid) in self.by_ident:
             try:
                 raise KeyError("uuid %r in use" % uuid)
             except:
                 content = error.wrap_exception()
                 self.log.error("uuid %r in use", uuid, exc_info=True)
         else:
             for h, ec in iteritems(self.incoming_registrations):
                 if uuid == h:
                     try:
                         raise KeyError("heart_id %r in use" % uuid)
                     except:
                         self.log.error("heart_id %r in use", uuid, exc_info=True)
                         content = error.wrap_exception()
                     break
                 elif uuid == ec.uuid:
                     try:
                         raise KeyError("uuid %r in use" % uuid)
                     except:
                         self.log.error("uuid %r in use", uuid, exc_info=True)
                         content = error.wrap_exception()
                     break
 
         msg = self.session.send(self.query, "registration_reply",
                 content=content,
                 ident=reg)
 
         heart = cast_bytes(uuid)
 
         if content['status'] == 'ok':
             if heart in self.heartmonitor.hearts:
                 # already beating
                 self.incoming_registrations[heart] = EngineConnector(id=eid,uuid=uuid)
                 self.finish_registration(heart)
             else:
                 purge = lambda : self._purge_stalled_registration(heart)
                 dc = ioloop.DelayedCallback(purge, self.registration_timeout, self.loop)
                 dc.start()
                 self.incoming_registrations[heart] = EngineConnector(id=eid,uuid=uuid,stallback=dc)
         else:
             self.log.error("registration::registration %i failed: %r", eid, content['evalue'])
         
         return eid
 
     def unregister_engine(self, ident, msg):
         """Unregister an engine that explicitly requested to leave."""
         try:
             eid = msg['content']['id']
         except:
             self.log.error("registration::bad engine id for unregistration: %r", ident, exc_info=True)
             return
         self.log.info("registration::unregister_engine(%r)", eid)
         # print (eid)
         uuid = self.keytable[eid]
         content=dict(id=eid, uuid=uuid)
         self.dead_engines.add(uuid)
         # self.ids.remove(eid)
         # uuid = self.keytable.pop(eid)
         #
         # ec = self.engines.pop(eid)
         # self.hearts.pop(ec.heartbeat)
         # self.by_ident.pop(ec.queue)
         # self.completed.pop(eid)
         handleit = lambda : self._handle_stranded_msgs(eid, uuid)
         dc = ioloop.DelayedCallback(handleit, self.registration_timeout, self.loop)
         dc.start()
         ############## TODO: HANDLE IT ################
         
         self._save_engine_state()
 
         if self.notifier:
             self.session.send(self.notifier, "unregistration_notification", content=content)
 
     def _handle_stranded_msgs(self, eid, uuid):
         """Handle messages known to be on an engine when the engine unregisters.
 
         It is possible that this will fire prematurely - that is, an engine will
         go down after completing a result, and the client will be notified
         that the result failed and later receive the actual result.
         """
 
         outstanding = self.queues[eid]
 
         for msg_id in outstanding:
             self.pending.remove(msg_id)
             self.all_completed.add(msg_id)
             try:
                 raise error.EngineError("Engine %r died while running task %r" % (eid, msg_id))
             except:
                 content = error.wrap_exception()
             # build a fake header:
             header = {}
             header['engine'] = uuid
             header['date'] = datetime.now()
             rec = dict(result_content=content, result_header=header, result_buffers=[])
             rec['completed'] = header['date']
             rec['engine_uuid'] = uuid
             try:
                 self.db.update_record(msg_id, rec)
             except Exception:
                 self.log.error("DB Error handling stranded msg %r", msg_id, exc_info=True)
 
 
     def finish_registration(self, heart):
         """Second half of engine registration, called after our HeartMonitor
         has received a beat from the Engine's Heart."""
         try:
             ec = self.incoming_registrations.pop(heart)
         except KeyError:
             self.log.error("registration::tried to finish nonexistant registration", exc_info=True)
             return
         self.log.info("registration::finished registering engine %i:%s", ec.id, ec.uuid)
         if ec.stallback is not None:
             ec.stallback.stop()
         eid = ec.id
         self.ids.add(eid)
         self.keytable[eid] = ec.uuid
         self.engines[eid] = ec
         self.by_ident[cast_bytes(ec.uuid)] = ec.id
         self.queues[eid] = list()
         self.tasks[eid] = list()
         self.completed[eid] = list()
         self.hearts[heart] = eid
         content = dict(id=eid, uuid=self.engines[eid].uuid)
         if self.notifier:
             self.session.send(self.notifier, "registration_notification", content=content)
         self.log.info("engine::Engine Connected: %i", eid)
         
         self._save_engine_state()
 
     def _purge_stalled_registration(self, heart):
         if heart in self.incoming_registrations:
             ec = self.incoming_registrations.pop(heart)
             self.log.info("registration::purging stalled registration: %i", ec.id)
         else:
             pass
 
     #-------------------------------------------------------------------------
     # Engine State
     #-------------------------------------------------------------------------
 
 
     def _cleanup_engine_state_file(self):
         """cleanup engine state mapping"""
         
         if os.path.exists(self.engine_state_file):
             self.log.debug("cleaning up engine state: %s", self.engine_state_file)
             try:
                 os.remove(self.engine_state_file)
             except IOError:
                 self.log.error("Couldn't cleanup file: %s", self.engine_state_file, exc_info=True)
 
 
     def _save_engine_state(self):
         """save engine mapping to JSON file"""
         if not self.engine_state_file:
             return
         self.log.debug("save engine state to %s" % self.engine_state_file)
         state = {}
         engines = {}
         for eid, ec in iteritems(self.engines):
             if ec.uuid not in self.dead_engines:
                 engines[eid] = ec.uuid
         
         state['engines'] = engines
         
         state['next_id'] = self._idcounter
         
         with open(self.engine_state_file, 'w') as f:
             json.dump(state, f)
 
 
     def _load_engine_state(self):
         """load engine mapping from JSON file"""
         if not os.path.exists(self.engine_state_file):
             return
         
         self.log.info("loading engine state from %s" % self.engine_state_file)
         
         with open(self.engine_state_file) as f:
             state = json.load(f)
         
         save_notifier = self.notifier
         self.notifier = None
         for eid, uuid in iteritems(state['engines']):
             heart = uuid.encode('ascii')
             # start with this heart as current and beating:
             self.heartmonitor.responses.add(heart)
             self.heartmonitor.hearts.add(heart)
             
             self.incoming_registrations[heart] = EngineConnector(id=int(eid), uuid=uuid)
             self.finish_registration(heart)
         
         self.notifier = save_notifier
         
         self._idcounter = state['next_id']
 
     #-------------------------------------------------------------------------
     # Client Requests
     #-------------------------------------------------------------------------
 
     def shutdown_request(self, client_id, msg):
         """handle shutdown request."""
         self.session.send(self.query, 'shutdown_reply', content={'status': 'ok'}, ident=client_id)
         # also notify other clients of shutdown
         self.session.send(self.notifier, 'shutdown_notice', content={'status': 'ok'})
         dc = ioloop.DelayedCallback(lambda : self._shutdown(), 1000, self.loop)
         dc.start()
 
     def _shutdown(self):
         self.log.info("hub::hub shutting down.")
         time.sleep(0.1)
         sys.exit(0)
 
 
     def check_load(self, client_id, msg):
         content = msg['content']
         try:
             targets = content['targets']
             targets = self._validate_targets(targets)
         except:
             content = error.wrap_exception()
             self.session.send(self.query, "hub_error",
                     content=content, ident=client_id)
             return
 
         content = dict(status='ok')
         # loads = {}
         for t in targets:
             content[bytes(t)] = len(self.queues[t])+len(self.tasks[t])
         self.session.send(self.query, "load_reply", content=content, ident=client_id)
 
 
     def queue_status(self, client_id, msg):
         """Return the Queue status of one or more targets.
 
         If verbose, return the msg_ids, else return len of each type.
 
         Keys:
 
         * queue (pending MUX jobs)
         * tasks (pending Task jobs)
         * completed (finished jobs from both queues)
         """
         content = msg['content']
         targets = content['targets']
         try:
             targets = self._validate_targets(targets)
         except:
             content = error.wrap_exception()
             self.session.send(self.query, "hub_error",
                     content=content, ident=client_id)
             return
         verbose = content.get('verbose', False)
         content = dict(status='ok')
         for t in targets:
             queue = self.queues[t]
             completed = self.completed[t]
             tasks = self.tasks[t]
             if not verbose:
                 queue = len(queue)
                 completed = len(completed)
                 tasks = len(tasks)
             content[str(t)] = {'queue': queue, 'completed': completed , 'tasks': tasks}
         content['unassigned'] = list(self.unassigned) if verbose else len(self.unassigned)
         # print (content)
         self.session.send(self.query, "queue_reply", content=content, ident=client_id)
 
     def purge_results(self, client_id, msg):
         """Purge results from memory. This method is more valuable before we move
         to a DB based message storage mechanism."""
         content = msg['content']
         self.log.info("Dropping records with %s", content)
         msg_ids = content.get('msg_ids', [])
         reply = dict(status='ok')
         if msg_ids == 'all':
             try:
                 self.db.drop_matching_records(dict(completed={'$ne':None}))
             except Exception:
                 reply = error.wrap_exception()
+                self.log.exception("Error dropping records")
         else:
             pending = [m for m in msg_ids if (m in self.pending)]
             if pending:
                 try:
                     raise IndexError("msg pending: %r" % pending[0])
                 except:
                     reply = error.wrap_exception()
+                    self.log.exception("Error dropping records")
             else:
                 try:
                     self.db.drop_matching_records(dict(msg_id={'$in':msg_ids}))
                 except Exception:
                     reply = error.wrap_exception()
+                    self.log.exception("Error dropping records")
 
             if reply['status'] == 'ok':
                 eids = content.get('engine_ids', [])
                 for eid in eids:
                     if eid not in self.engines:
                         try:
                             raise IndexError("No such engine: %i" % eid)
                         except:
                             reply = error.wrap_exception()
+                            self.log.exception("Error dropping records")
                         break
                     uid = self.engines[eid].uuid
                     try:
                         self.db.drop_matching_records(dict(engine_uuid=uid, completed={'$ne':None}))
                     except Exception:
                         reply = error.wrap_exception()
+                        self.log.exception("Error dropping records")
                         break
 
         self.session.send(self.query, 'purge_reply', content=reply, ident=client_id)
 
     def resubmit_task(self, client_id, msg):
         """Resubmit one or more tasks."""
         def finish(reply):
             self.session.send(self.query, 'resubmit_reply', content=reply, ident=client_id)
 
         content = msg['content']
         msg_ids = content['msg_ids']
         reply = dict(status='ok')
         try:
             records = self.db.find_records({'msg_id' : {'$in' : msg_ids}}, keys=[
                 'header', 'content', 'buffers'])
         except Exception:
             self.log.error('db::db error finding tasks to resubmit', exc_info=True)
             return finish(error.wrap_exception())
 
         # validate msg_ids
         found_ids = [ rec['msg_id'] for rec in records ]
         pending_ids = [ msg_id for msg_id in found_ids if msg_id in self.pending ]
         if len(records) > len(msg_ids):
             try:
                 raise RuntimeError("DB appears to be in an inconsistent state."
                     "More matching records were found than should exist")
             except Exception:
+                self.log.exception("Failed to resubmit task")
                 return finish(error.wrap_exception())
         elif len(records) < len(msg_ids):
             missing = [ m for m in msg_ids if m not in found_ids ]
             try:
                 raise KeyError("No such msg(s): %r" % missing)
             except KeyError:
+                self.log.exception("Failed to resubmit task")
                 return finish(error.wrap_exception())
         elif pending_ids:
             pass
             # no need to raise on resubmit of pending task, now that we
             # resubmit under new ID, but do we want to raise anyway?
             # msg_id = invalid_ids[0]
             # try:
             #     raise ValueError("Task(s) %r appears to be inflight" % )
             # except Exception:
             #     return finish(error.wrap_exception())
 
         # mapping of original IDs to resubmitted IDs
         resubmitted = {}
 
         # send the messages
         for rec in records:
             header = rec['header']
             msg = self.session.msg(header['msg_type'], parent=header)
             msg_id = msg['msg_id']
             msg['content'] = rec['content']
             
             # use the old header, but update msg_id and timestamp
             fresh = msg['header']
             header['msg_id'] = fresh['msg_id']
             header['date'] = fresh['date']
             msg['header'] = header
 
             self.session.send(self.resubmit, msg, buffers=rec['buffers'])
 
             resubmitted[rec['msg_id']] = msg_id
             self.pending.add(msg_id)
             msg['buffers'] = rec['buffers']
             try:
                 self.db.add_record(msg_id, init_record(msg))
             except Exception:
                 self.log.error("db::DB Error updating record: %s", msg_id, exc_info=True)
                 return finish(error.wrap_exception())
 
         finish(dict(status='ok', resubmitted=resubmitted))
         
         # store the new IDs in the Task DB
         for msg_id, resubmit_id in iteritems(resubmitted):
             try:
                 self.db.update_record(msg_id, {'resubmitted' : resubmit_id})
             except Exception:
                 self.log.error("db::DB Error updating record: %s", msg_id, exc_info=True)
 
 
     def _extract_record(self, rec):
         """decompose a TaskRecord dict into subsection of reply for get_result"""
         io_dict = {}
         for key in ('pyin', 'pyout', 'pyerr', 'stdout', 'stderr'):
                 io_dict[key] = rec[key]
         content = { 
             'header': rec['header'],
             'metadata': rec['metadata'],
             'result_metadata': rec['result_metadata'],
             'result_header' : rec['result_header'],
             'result_content': rec['result_content'],
             'received' : rec['received'],
             'io' : io_dict,
         }
         if rec['result_buffers']:
             buffers = list(map(bytes, rec['result_buffers']))
         else:
             buffers = []
 
         return content, buffers
 
     def get_results(self, client_id, msg):
         """Get the result of 1 or more messages."""
         content = msg['content']
         msg_ids = sorted(set(content['msg_ids']))
         statusonly = content.get('status_only', False)
         pending = []
         completed = []
         content = dict(status='ok')
         content['pending'] = pending
         content['completed'] = completed
         buffers = []
         if not statusonly:
             try:
                 matches = self.db.find_records(dict(msg_id={'$in':msg_ids}))
                 # turn match list into dict, for faster lookup
                 records = {}
                 for rec in matches:
                     records[rec['msg_id']] = rec
             except Exception:
                 content = error.wrap_exception()
+                self.log.exception("Failed to get results")
                 self.session.send(self.query, "result_reply", content=content,
                                                     parent=msg, ident=client_id)
                 return
         else:
             records = {}
         for msg_id in msg_ids:
             if msg_id in self.pending:
                 pending.append(msg_id)
             elif msg_id in self.all_completed:
                 completed.append(msg_id)
                 if not statusonly:
                     c,bufs = self._extract_record(records[msg_id])
                     content[msg_id] = c
                     buffers.extend(bufs)
             elif msg_id in records:
                 if rec['completed']:
                     completed.append(msg_id)
                     c,bufs = self._extract_record(records[msg_id])
                     content[msg_id] = c
                     buffers.extend(bufs)
                 else:
                     pending.append(msg_id)
             else:
                 try:
                     raise KeyError('No such message: '+msg_id)
                 except:
                     content = error.wrap_exception()
                 break
         self.session.send(self.query, "result_reply", content=content,
                                             parent=msg, ident=client_id,
                                             buffers=buffers)
 
     def get_history(self, client_id, msg):
         """Get a list of all msg_ids in our DB records"""
         try:
             msg_ids = self.db.get_history()
         except Exception as e:
             content = error.wrap_exception()
+            self.log.exception("Failed to get history")
         else:
             content = dict(status='ok', history=msg_ids)
 
         self.session.send(self.query, "history_reply", content=content,
                                             parent=msg, ident=client_id)
 
     def db_query(self, client_id, msg):
         """Perform a raw query on the task record database."""
         content = msg['content']
         query = extract_dates(content.get('query', {}))
         keys = content.get('keys', None)
         buffers = []
         empty = list()
         try:
             records = self.db.find_records(query, keys)
         except Exception as e:
             content = error.wrap_exception()
+            self.log.exception("DB query failed")
         else:
             # extract buffers from reply content:
             if keys is not None:
                 buffer_lens = [] if 'buffers' in keys else None
                 result_buffer_lens = [] if 'result_buffers' in keys else None
             else:
                 buffer_lens = None
                 result_buffer_lens = None
 
             for rec in records:
                 # buffers may be None, so double check
                 b = rec.pop('buffers', empty) or empty
                 if buffer_lens is not None:
                     buffer_lens.append(len(b))
                     buffers.extend(b)
                 rb = rec.pop('result_buffers', empty) or empty
                 if result_buffer_lens is not None:
                     result_buffer_lens.append(len(rb))
                     buffers.extend(rb)
             content = dict(status='ok', records=records, buffer_lens=buffer_lens,
                                     result_buffer_lens=result_buffer_lens)
         # self.log.debug (content)
         self.session.send(self.query, "db_reply", content=content,
                                             parent=msg, ident=client_id,
                                             buffers=buffers)
 

IPython/parallel/tests/__init__.py

@@ -1,130 +1,145 b''
 """toplevel setup/teardown for parallel tests."""
 from __future__ import print_function
 
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # Imports
 #-------------------------------------------------------------------------------
 
 import os
 import tempfile
 import time
 from subprocess import Popen, PIPE, STDOUT
 
 import nose
 
 from IPython.utils.path import get_ipython_dir
-from IPython.parallel import Client
+from IPython.parallel import Client, error
 from IPython.parallel.apps.launcher import (LocalProcessLauncher,
                                                   ipengine_cmd_argv,
                                                   ipcontroller_cmd_argv,
                                                   SIGKILL,
                                                   ProcessStateError,
 )
 
 # globals
 launchers = []
 blackhole = open(os.devnull, 'w')
 
 # Launcher class
 class TestProcessLauncher(LocalProcessLauncher):
     """subclass LocalProcessLauncher, to prevent extra sockets and threads being created on Windows"""
     def start(self):
         if self.state == 'before':
             # Store stdout & stderr to show with failing tests.
             # This is defined in IPython.testing.iptest
             self.process = Popen(self.args,
                 stdout=nose.iptest_stdstreams_fileno(), stderr=STDOUT,
                 env=os.environ,
                 cwd=self.work_dir
             )
             self.notify_start(self.process.pid)
             self.poll = self.process.poll
         else:
             s = 'The process was already started and has state: %r' % self.state
             raise ProcessStateError(s)
 
 # nose setup/teardown
 
 def setup():
+    
+    # show tracebacks for RemoteErrors
+    class RemoteErrorWithTB(error.RemoteError):
+        def __str__(self):
+            s = super(RemoteErrorWithTB, self).__str__()
+            return '\n'.join([s, self.traceback or ''])
+    
+    error.RemoteError = RemoteErrorWithTB
+    
     cluster_dir = os.path.join(get_ipython_dir(), 'profile_iptest')
     engine_json = os.path.join(cluster_dir, 'security', 'ipcontroller-engine.json')
     client_json = os.path.join(cluster_dir, 'security', 'ipcontroller-client.json')
     for json in (engine_json, client_json):
         if os.path.exists(json):
             os.remove(json)
     
     cp = TestProcessLauncher()
     cp.cmd_and_args = ipcontroller_cmd_argv + \
                 ['--profile=iptest', '--log-level=20', '--ping=250', '--dictdb']
     cp.start()
     launchers.append(cp)
     tic = time.time()
     while not os.path.exists(engine_json) or not os.path.exists(client_json):
         if cp.poll() is not None:
             raise RuntimeError("The test controller exited with status %s" % cp.poll())
         elif time.time()-tic > 15:
             raise RuntimeError("Timeout waiting for the test controller to start.")
         time.sleep(0.1)
     add_engines(1)
 
 def add_engines(n=1, profile='iptest', total=False):
     """add a number of engines to a given profile.
     
     If total is True, then already running engines are counted, and only
     the additional engines necessary (if any) are started.
     """
     rc = Client(profile=profile)
     base = len(rc)
     
     if total:
         n = max(n - base, 0)
     
     eps = []
     for i in range(n):
         ep = TestProcessLauncher()
         ep.cmd_and_args = ipengine_cmd_argv + [
             '--profile=%s' % profile,
             '--log-level=50',
             '--InteractiveShell.colors=nocolor'
             ]
         ep.start()
         launchers.append(ep)
         eps.append(ep)
     tic = time.time()
     while len(rc) < base+n:
         if any([ ep.poll() is not None for ep in eps ]):
             raise RuntimeError("A test engine failed to start.")
         elif time.time()-tic > 15:
             raise RuntimeError("Timeout waiting for engines to connect.")
         time.sleep(.1)
         rc.spin()
     rc.close()
     return eps
 
 def teardown():
-    time.sleep(1)
+    try:
+        time.sleep(1)
+    except KeyboardInterrupt:
+        return
     while launchers:
         p = launchers.pop()
         if p.poll() is None:
             try:
                 p.stop()
             except Exception as e:
                 print(e)
                 pass
         if p.poll() is None:
-            time.sleep(.25)
+            try:
+                time.sleep(.25)
+            except KeyboardInterrupt:
+                return
         if p.poll() is None:
             try:
                 print('cleaning up test process...')
                 p.signal(SIGKILL)
             except:
                 print("couldn't shutdown process: ", p)
     blackhole.close()
     

IPython/parallel/tests/test_db.py

@@ -1,314 +1,314 b''
 """Tests for db backends
 
 Authors:
 
 * Min RK
 """
 
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # Imports
 #-------------------------------------------------------------------------------
 
 from __future__ import division
 
 import logging
 import os
 import tempfile
 import time
 
 from datetime import datetime, timedelta
 from unittest import TestCase
 
 from IPython.parallel import error
 from IPython.parallel.controller.dictdb import DictDB
 from IPython.parallel.controller.sqlitedb import SQLiteDB
 from IPython.parallel.controller.hub import init_record, empty_record
 
 from IPython.testing import decorators as dec
 from IPython.kernel.zmq.session import Session
 
 
 #-------------------------------------------------------------------------------
 # TestCases
 #-------------------------------------------------------------------------------
 
 
 def setup():
     global temp_db
     temp_db = tempfile.NamedTemporaryFile(suffix='.db').name
 
 
 class TaskDBTest:
     def setUp(self):
         self.session = Session()
         self.db = self.create_db()
         self.load_records(16)
     
     def create_db(self):
         raise NotImplementedError
     
     def load_records(self, n=1, buffer_size=100):
         """load n records for testing"""
         #sleep 1/10 s, to ensure timestamp is different to previous calls
         time.sleep(0.1)
         msg_ids = []
         for i in range(n):
             msg = self.session.msg('apply_request', content=dict(a=5))
             msg['buffers'] = [os.urandom(buffer_size)]
             rec = init_record(msg)
             msg_id = msg['header']['msg_id']
             msg_ids.append(msg_id)
             self.db.add_record(msg_id, rec)
         return msg_ids
     
     def test_add_record(self):
         before = self.db.get_history()
         self.load_records(5)
         after = self.db.get_history()
         self.assertEqual(len(after), len(before)+5)
         self.assertEqual(after[:-5],before)
         
     def test_drop_record(self):
         msg_id = self.load_records()[-1]
         rec = self.db.get_record(msg_id)
         self.db.drop_record(msg_id)
         self.assertRaises(KeyError,self.db.get_record, msg_id)
     
     def _round_to_millisecond(self, dt):
         """necessary because mongodb rounds microseconds"""
         micro = dt.microsecond
         extra = int(str(micro)[-3:])
         return dt - timedelta(microseconds=extra)
     
     def test_update_record(self):
         now = self._round_to_millisecond(datetime.now())
         # 
         msg_id = self.db.get_history()[-1]
         rec1 = self.db.get_record(msg_id)
         data = {'stdout': 'hello there', 'completed' : now}
         self.db.update_record(msg_id, data)
         rec2 = self.db.get_record(msg_id)
         self.assertEqual(rec2['stdout'], 'hello there')
         self.assertEqual(rec2['completed'], now)
         rec1.update(data)
         self.assertEqual(rec1, rec2)
     
     # def test_update_record_bad(self):
     #     """test updating nonexistant records"""
     #     msg_id = str(uuid.uuid4())
     #     data = {'stdout': 'hello there'}
     #     self.assertRaises(KeyError, self.db.update_record, msg_id, data)
 
     def test_find_records_dt(self):
         """test finding records by date"""
         hist = self.db.get_history()
         middle = self.db.get_record(hist[len(hist)//2])
         tic = middle['submitted']
         before = self.db.find_records({'submitted' : {'$lt' : tic}})
         after = self.db.find_records({'submitted' : {'$gte' : tic}})
         self.assertEqual(len(before)+len(after),len(hist))
         for b in before:
             self.assertTrue(b['submitted'] < tic)
         for a in after:
             self.assertTrue(a['submitted'] >= tic)
         same = self.db.find_records({'submitted' : tic})
         for s in same:
             self.assertTrue(s['submitted'] == tic)
     
     def test_find_records_keys(self):
         """test extracting subset of record keys"""
         found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['submitted', 'completed'])
         for rec in found:
             self.assertEqual(set(rec.keys()), set(['msg_id', 'submitted', 'completed']))
     
     def test_find_records_msg_id(self):
         """ensure msg_id is always in found records"""
         found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['submitted', 'completed'])
         for rec in found:
             self.assertTrue('msg_id' in rec.keys())
         found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['submitted'])
         for rec in found:
             self.assertTrue('msg_id' in rec.keys())
         found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['msg_id'])
         for rec in found:
             self.assertTrue('msg_id' in rec.keys())
     
     def test_find_records_in(self):
         """test finding records with '$in','$nin' operators"""
         hist = self.db.get_history()
         even = hist[::2]
         odd = hist[1::2]
         recs = self.db.find_records({ 'msg_id' : {'$in' : even}})
         found = [ r['msg_id'] for r in recs ]
         self.assertEqual(set(even), set(found))
         recs = self.db.find_records({ 'msg_id' : {'$nin' : even}})
         found = [ r['msg_id'] for r in recs ]
         self.assertEqual(set(odd), set(found))
     
     def test_get_history(self):
         msg_ids = self.db.get_history()
         latest = datetime(1984,1,1)
         for msg_id in msg_ids:
             rec = self.db.get_record(msg_id)
             newt = rec['submitted']
             self.assertTrue(newt >= latest)
             latest = newt
         msg_id = self.load_records(1)[-1]
         self.assertEqual(self.db.get_history()[-1],msg_id)
     
     def test_datetime(self):
         """get/set timestamps with datetime objects"""
         msg_id = self.db.get_history()[-1]
         rec = self.db.get_record(msg_id)
         self.assertTrue(isinstance(rec['submitted'], datetime))
         self.db.update_record(msg_id, dict(completed=datetime.now()))
         rec = self.db.get_record(msg_id)
         self.assertTrue(isinstance(rec['completed'], datetime))
 
     def test_drop_matching(self):
         msg_ids = self.load_records(10)
         query = {'msg_id' : {'$in':msg_ids}}
         self.db.drop_matching_records(query)
         recs = self.db.find_records(query)
         self.assertEqual(len(recs), 0)
     
     def test_null(self):
         """test None comparison queries"""
         msg_ids = self.load_records(10)
 
         query = {'msg_id' : None}
         recs = self.db.find_records(query)
         self.assertEqual(len(recs), 0)
 
         query = {'msg_id' : {'$ne' : None}}
         recs = self.db.find_records(query)
         self.assertTrue(len(recs) >= 10)
     
     def test_pop_safe_get(self):
         """editing query results shouldn't affect record [get]"""
         msg_id = self.db.get_history()[-1]
         rec = self.db.get_record(msg_id)
         rec.pop('buffers')
         rec['garbage'] = 'hello'
         rec['header']['msg_id'] = 'fubar'
         rec2 = self.db.get_record(msg_id)
         self.assertTrue('buffers' in rec2)
         self.assertFalse('garbage' in rec2)
         self.assertEqual(rec2['header']['msg_id'], msg_id)
     
     def test_pop_safe_find(self):
         """editing query results shouldn't affect record [find]"""
         msg_id = self.db.get_history()[-1]
         rec = self.db.find_records({'msg_id' : msg_id})[0]
         rec.pop('buffers')
         rec['garbage'] = 'hello'
         rec['header']['msg_id'] = 'fubar'
         rec2 = self.db.find_records({'msg_id' : msg_id})[0]
         self.assertTrue('buffers' in rec2)
         self.assertFalse('garbage' in rec2)
         self.assertEqual(rec2['header']['msg_id'], msg_id)
 
     def test_pop_safe_find_keys(self):
         """editing query results shouldn't affect record [find+keys]"""
         msg_id = self.db.get_history()[-1]
         rec = self.db.find_records({'msg_id' : msg_id}, keys=['buffers', 'header'])[0]
         rec.pop('buffers')
         rec['garbage'] = 'hello'
         rec['header']['msg_id'] = 'fubar'
         rec2 = self.db.find_records({'msg_id' : msg_id})[0]
         self.assertTrue('buffers' in rec2)
         self.assertFalse('garbage' in rec2)
         self.assertEqual(rec2['header']['msg_id'], msg_id)
 
 
 class TestDictBackend(TaskDBTest, TestCase):
     
     def create_db(self):
         return DictDB()
     
     def test_cull_count(self):
         self.db = self.create_db() # skip the load-records init from setUp
         self.db.record_limit = 20
         self.db.cull_fraction = 0.2
         self.load_records(20)
         self.assertEqual(len(self.db.get_history()), 20)
         self.load_records(1)
         # 0.2 * 20 = 4, 21 - 4 = 17
         self.assertEqual(len(self.db.get_history()), 17)
         self.load_records(3)
         self.assertEqual(len(self.db.get_history()), 20)
         self.load_records(1)
         self.assertEqual(len(self.db.get_history()), 17)
         
-        for i in range(100):
+        for i in range(25):
             self.load_records(1)
             self.assertTrue(len(self.db.get_history()) >= 17)
             self.assertTrue(len(self.db.get_history()) <= 20)
 
     def test_cull_size(self):
         self.db = self.create_db() # skip the load-records init from setUp
         self.db.size_limit = 1000
         self.db.cull_fraction = 0.2
         self.load_records(100, buffer_size=10)
         self.assertEqual(len(self.db.get_history()), 100)
         self.load_records(1, buffer_size=0)
         self.assertEqual(len(self.db.get_history()), 101)
         self.load_records(1, buffer_size=1)
         # 0.2 * 100 = 20, 101 - 20 = 81
         self.assertEqual(len(self.db.get_history()), 81)
     
     def test_cull_size_drop(self):
         """dropping records updates tracked buffer size"""
         self.db = self.create_db() # skip the load-records init from setUp
         self.db.size_limit = 1000
         self.db.cull_fraction = 0.2
         self.load_records(100, buffer_size=10)
         self.assertEqual(len(self.db.get_history()), 100)
         self.db.drop_record(self.db.get_history()[-1])
         self.assertEqual(len(self.db.get_history()), 99)
         self.load_records(1, buffer_size=5)
         self.assertEqual(len(self.db.get_history()), 100)
         self.load_records(1, buffer_size=5)
         self.assertEqual(len(self.db.get_history()), 101)
         self.load_records(1, buffer_size=1)
         self.assertEqual(len(self.db.get_history()), 81)
 
     def test_cull_size_update(self):
         """updating records updates tracked buffer size"""
         self.db = self.create_db() # skip the load-records init from setUp
         self.db.size_limit = 1000
         self.db.cull_fraction = 0.2
         self.load_records(100, buffer_size=10)
         self.assertEqual(len(self.db.get_history()), 100)
         msg_id = self.db.get_history()[-1]
         self.db.update_record(msg_id, dict(result_buffers = [os.urandom(10)], buffers=[]))
         self.assertEqual(len(self.db.get_history()), 100)
         self.db.update_record(msg_id, dict(result_buffers = [os.urandom(11)], buffers=[]))
         self.assertEqual(len(self.db.get_history()), 79)
 
 class TestSQLiteBackend(TaskDBTest, TestCase):
 
     @dec.skip_without('sqlite3')
     def create_db(self):
         location, fname = os.path.split(temp_db)
         log = logging.getLogger('test')
         log.setLevel(logging.CRITICAL)
         return SQLiteDB(location=location, fname=fname, log=log)
     
     def tearDown(self):
         self.db._db.close()
 
 
 def teardown():
     """cleanup task db file after all tests have run"""
     try:
         os.remove(temp_db)
     except:
         pass

IPython/utils/jsonutil.py

@@ -1,256 +1,259 b''
 """Utilities to manipulate JSON objects.
 """
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2010-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING.txt, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 # stdlib
 import math
 import re
 import types
 from datetime import datetime
 
 try:
     # base64.encodestring is deprecated in Python 3.x
     from base64 import encodebytes
 except ImportError:
     # Python 2.x
     from base64 import encodestring as encodebytes
 
 from IPython.utils import py3compat
 from IPython.utils.py3compat import string_types, unicode_type, iteritems
 from IPython.utils.encoding import DEFAULT_ENCODING
 next_attr_name = '__next__' if py3compat.PY3 else 'next'
 
 #-----------------------------------------------------------------------------
 # Globals and constants
 #-----------------------------------------------------------------------------
 
 # timestamp formats
 ISO8601 = "%Y-%m-%dT%H:%M:%S.%f"
-ISO8601_PAT=re.compile(r"^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{1,6})Z?([\+\-]\d{2}:?\d{2})?$")
+ISO8601_PAT=re.compile(r"^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(\.\d{1,6})?Z?([\+\-]\d{2}:?\d{2})?$")
 
 #-----------------------------------------------------------------------------
 # Classes and functions
 #-----------------------------------------------------------------------------
 
 def rekey(dikt):
     """Rekey a dict that has been forced to use str keys where there should be
     ints by json."""
     for k in dikt:
         if isinstance(k, string_types):
             ik=fk=None
             try:
                 ik = int(k)
             except ValueError:
                 try:
                     fk = float(k)
                 except ValueError:
                     continue
             if ik is not None:
                 nk = ik
             else:
                 nk = fk
             if nk in dikt:
                 raise KeyError("already have key %r"%nk)
             dikt[nk] = dikt.pop(k)
     return dikt
 
 def parse_date(s):
     """parse an ISO8601 date string
     
     If it is None or not a valid ISO8601 timestamp,
     it will be returned unmodified.
     Otherwise, it will return a datetime object.
     """
     if s is None:
         return s
     m = ISO8601_PAT.match(s)
     if m:
         # FIXME: add actual timezone support
         # this just drops the timezone info
-        notz = m.groups()[0]
+        notz, ms, tz = m.groups()
+        if not ms:
+            ms = '.0'
+        notz = notz + ms
         return datetime.strptime(notz, ISO8601)
     return s
 
 def extract_dates(obj):
     """extract ISO8601 dates from unpacked JSON"""
     if isinstance(obj, dict):
         new_obj = {} # don't clobber
         for k,v in iteritems(obj):
             new_obj[k] = extract_dates(v)
         obj = new_obj
     elif isinstance(obj, (list, tuple)):
         obj = [ extract_dates(o) for o in obj ]
     elif isinstance(obj, string_types):
         obj = parse_date(obj)
     return obj
 
 def squash_dates(obj):
     """squash datetime objects into ISO8601 strings"""
     if isinstance(obj, dict):
         obj = dict(obj) # don't clobber
         for k,v in iteritems(obj):
             obj[k] = squash_dates(v)
     elif isinstance(obj, (list, tuple)):
         obj = [ squash_dates(o) for o in obj ]
     elif isinstance(obj, datetime):
         obj = obj.isoformat()
     return obj
 
 def date_default(obj):
     """default function for packing datetime objects in JSON."""
     if isinstance(obj, datetime):
         return obj.isoformat()
     else:
         raise TypeError("%r is not JSON serializable"%obj)
 
 
 # constants for identifying png/jpeg data
 PNG = b'\x89PNG\r\n\x1a\n'
 # front of PNG base64-encoded
 PNG64 = b'iVBORw0KG'
 JPEG = b'\xff\xd8'
 # front of JPEG base64-encoded
 JPEG64 = b'/9'
 # front of PDF base64-encoded
 PDF64 = b'JVBER'
 
 def encode_images(format_dict):
     """b64-encodes images in a displaypub format dict
 
     Perhaps this should be handled in json_clean itself?
 
     Parameters
     ----------
 
     format_dict : dict
         A dictionary of display data keyed by mime-type
 
     Returns
     -------
 
     format_dict : dict
         A copy of the same dictionary,
         but binary image data ('image/png', 'image/jpeg' or 'application/pdf')
         is base64-encoded.
 
     """
     encoded = format_dict.copy()
 
     pngdata = format_dict.get('image/png')
     if isinstance(pngdata, bytes):
         # make sure we don't double-encode
         if not pngdata.startswith(PNG64):
             pngdata = encodebytes(pngdata)
         encoded['image/png'] = pngdata.decode('ascii')
 
     jpegdata = format_dict.get('image/jpeg')
     if isinstance(jpegdata, bytes):
         # make sure we don't double-encode
         if not jpegdata.startswith(JPEG64):
             jpegdata = encodebytes(jpegdata)
         encoded['image/jpeg'] = jpegdata.decode('ascii')
 
     pdfdata = format_dict.get('application/pdf')
     if isinstance(pdfdata, bytes):
         # make sure we don't double-encode
         if not pdfdata.startswith(PDF64):
             pdfdata = encodebytes(pdfdata)
         encoded['application/pdf'] = pdfdata.decode('ascii')
 
     return encoded
 
 
 def json_clean(obj):
     """Clean an object to ensure it's safe to encode in JSON.
 
     Atomic, immutable objects are returned unmodified.  Sets and tuples are
     converted to lists, lists are copied and dicts are also copied.
 
     Note: dicts whose keys could cause collisions upon encoding (such as a dict
     with both the number 1 and the string '1' as keys) will cause a ValueError
     to be raised.
 
     Parameters
     ----------
     obj : any python object
 
     Returns
     -------
     out : object
 
       A version of the input which will not cause an encoding error when
       encoded as JSON.  Note that this function does not *encode* its inputs,
       it simply sanitizes it so that there will be no encoding errors later.
 
     Examples
     --------
     >>> json_clean(4)
     4
     >>> json_clean(list(range(10)))
     [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
     >>> sorted(json_clean(dict(x=1, y=2)).items())
     [('x', 1), ('y', 2)]
     >>> sorted(json_clean(dict(x=1, y=2, z=[1,2,3])).items())
     [('x', 1), ('y', 2), ('z', [1, 2, 3])]
     >>> json_clean(True)
     True
     """
     # types that are 'atomic' and ok in json as-is.
     atomic_ok = (unicode_type, type(None))
 
     # containers that we need to convert into lists
     container_to_list = (tuple, set, types.GeneratorType)
 
     if isinstance(obj, float):
         # cast out-of-range floats to their reprs
         if math.isnan(obj) or math.isinf(obj):
             return repr(obj)
         return float(obj)
     
     if isinstance(obj, int):
         # cast int to int, in case subclasses override __str__ (e.g. boost enum, #4598)
         if isinstance(obj, bool):
             # bools are ints, but we don't want to cast them to 0,1
             return obj
         return int(obj)
 
     if isinstance(obj, atomic_ok):
         return obj
 
     if isinstance(obj, bytes):
         return obj.decode(DEFAULT_ENCODING, 'replace')
 
     if isinstance(obj, container_to_list) or (
         hasattr(obj, '__iter__') and hasattr(obj, next_attr_name)):
         obj = list(obj)
 
     if isinstance(obj, list):
         return [json_clean(x) for x in obj]
 
     if isinstance(obj, dict):
         # First, validate that the dict won't lose data in conversion due to
         # key collisions after stringification.  This can happen with keys like
         # True and 'true' or 1 and '1', which collide in JSON.
         nkeys = len(obj)
         nkeys_collapsed = len(set(map(str, obj)))
         if nkeys != nkeys_collapsed:
             raise ValueError('dict can not be safely converted to JSON: '
                              'key collision would lead to dropped values')
         # If all OK, proceed by making the new dict that will be json-safe
         out = {}
         for k,v in iteritems(obj):
             out[str(k)] = json_clean(v)
         return out
 
     # If we get here, we don't know how to handle the object, so we just get
     # its repr and return that.  This will catch lambdas, open sockets, class
     # objects, and any other complicated contraption that json can't encode
     return repr(obj)

IPython/utils/tests/test_jsonutil.py

@@ -1,149 +1,151 b''
 """Test suite for our JSON utilities.
 """
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2010-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING.txt, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 # stdlib
 import datetime
 import json
 from base64 import decodestring
 
 # third party
 import nose.tools as nt
 
 # our own
 from IPython.utils import jsonutil, tz
 from ..jsonutil import json_clean, encode_images
 from ..py3compat import unicode_to_str, str_to_bytes, iteritems
 
 #-----------------------------------------------------------------------------
 # Test functions
 #-----------------------------------------------------------------------------
 class Int(int):
     def __str__(self):
         return 'Int(%i)' % self
 
 def test():
     # list of input/expected output.  Use None for the expected output if it
     # can be the same as the input.
     pairs = [(1, None), # start with scalars
              (1.0, None),
              ('a', None),
              (True, None),
              (False, None),
              (None, None),
              # complex numbers for now just go to strings, as otherwise they
              # are unserializable
              (1j, '1j'),
              # Containers
              ([1, 2], None),
              ((1, 2), [1, 2]),
              (set([1, 2]), [1, 2]),
              (dict(x=1), None),
              ({'x': 1, 'y':[1,2,3], '1':'int'}, None),
              # More exotic objects
              ((x for x in range(3)), [0, 1, 2]),
              (iter([1, 2]), [1, 2]),
              (Int(5), 5),
              ]
     
     for val, jval in pairs:
         if jval is None:
             jval = val
         out = json_clean(val)
         # validate our cleanup
         nt.assert_equal(out, jval)
         # and ensure that what we return, indeed encodes cleanly
         json.loads(json.dumps(out))
 
 
 
 def test_encode_images():
     # invalid data, but the header and footer are from real files
     pngdata = b'\x89PNG\r\n\x1a\nblahblahnotactuallyvalidIEND\xaeB`\x82'
     jpegdata = b'\xff\xd8\xff\xe0\x00\x10JFIFblahblahjpeg(\xa0\x0f\xff\xd9'
     pdfdata = b'%PDF-1.\ntrailer<</Root<</Pages<</Kids[<</MediaBox[0 0 3 3]>>]>>>>>>'
     
     fmt = {
         'image/png'  : pngdata,
         'image/jpeg' : jpegdata,
         'application/pdf' : pdfdata
     }
     encoded = encode_images(fmt)
     for key, value in iteritems(fmt):
         # encoded has unicode, want bytes
         decoded = decodestring(encoded[key].encode('ascii'))
         nt.assert_equal(decoded, value)
     encoded2 = encode_images(encoded)
     nt.assert_equal(encoded, encoded2)
     
     b64_str = {}
     for key, encoded in iteritems(encoded):
         b64_str[key] = unicode_to_str(encoded)
     encoded3 = encode_images(b64_str)
     nt.assert_equal(encoded3, b64_str)
     for key, value in iteritems(fmt):
         # encoded3 has str, want bytes
         decoded = decodestring(str_to_bytes(encoded3[key]))
         nt.assert_equal(decoded, value)
 
 def test_lambda():
     jc = json_clean(lambda : 1)
-    assert isinstance(jc, str)
-    assert '<lambda>' in jc
+    nt.assert_is_instance(jc, str)
+    nt.assert_in('<lambda>', jc)
     json.dumps(jc)
 
 def test_extract_dates():
     timestamps = [
         '2013-07-03T16:34:52.249482',
         '2013-07-03T16:34:52.249482Z',
         '2013-07-03T16:34:52.249482Z-0800',
         '2013-07-03T16:34:52.249482Z+0800',
         '2013-07-03T16:34:52.249482Z+08:00',
         '2013-07-03T16:34:52.249482Z-08:00',
         '2013-07-03T16:34:52.249482-0800',
         '2013-07-03T16:34:52.249482+0800',
         '2013-07-03T16:34:52.249482+08:00',
         '2013-07-03T16:34:52.249482-08:00',
     ]
     extracted = jsonutil.extract_dates(timestamps)
     ref = extracted[0]
     for dt in extracted:
         nt.assert_true(isinstance(dt, datetime.datetime))
         nt.assert_equal(dt, ref)
 
 def test_parse_ms_precision():
-    base = '2013-07-03T16:34:52.'
+    base = '2013-07-03T16:34:52'
     digits = '1234567890'
     
+    parsed = jsonutil.parse_date(base)
+    nt.assert_is_instance(parsed, datetime.datetime)
     for i in range(len(digits)):
-        ts = base + digits[:i]
+        ts = base + '.' + digits[:i]
         parsed = jsonutil.parse_date(ts)
         if i >= 1 and i <= 6:
-            assert isinstance(parsed, datetime.datetime)
+            nt.assert_is_instance(parsed, datetime.datetime)
         else:
-            assert isinstance(parsed, str)
+            nt.assert_is_instance(parsed, str)
 
 def test_date_default():
     data = dict(today=datetime.datetime.now(), utcnow=tz.utcnow())
     jsondata = json.dumps(data, default=jsonutil.date_default)
     nt.assert_in("+00", jsondata)
     nt.assert_equal(jsondata.count("+00"), 1)
     extracted = jsonutil.extract_dates(json.loads(jsondata))
     for dt in extracted.values():
-        nt.assert_true(isinstance(dt, datetime.datetime))
+        nt.assert_is_instance(dt, datetime.datetime)
 
 def test_exception():
     bad_dicts = [{1:'number', '1':'string'},
                  {True:'bool', 'True':'string'},
                  ]
     for d in bad_dicts:
         nt.assert_raises(ValueError, json_clean, d)