diff --git a/mitogen/core.py b/mitogen/core.py index dadf0924..d6134694 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -1279,7 +1279,7 @@ def _unpickle_context(router, context_id, name): (isinstance(name, UnicodeType) and len(name) < 100)) ): raise TypeError('cannot unpickle Context: bad input') - return router.context_class(router, context_id, name) + return router.context_by_id(context_id, name=name) class Poller(object): @@ -1732,6 +1732,15 @@ class Router(object): _, (_, func, _) = self._handle_map.popitem() func(Message.dead()) + def context_by_id(self, context_id, via_id=None, create=True, name=None): + context = self._context_by_id.get(context_id) + if create and not context: + context = self.context_class(self, context_id, name=name) + if via_id is not None: + context.via = self.context_by_id(via_id) + self._context_by_id[context_id] = context + return context + def register(self, context, stream): _v and LOG.debug('register(%r, %r)', context, stream) self._stream_by_id[context.context_id] = stream diff --git a/mitogen/parent.py b/mitogen/parent.py index a57ca20b..d8a1a633 100644 --- a/mitogen/parent.py +++ b/mitogen/parent.py @@ -1581,15 +1581,6 @@ class Router(mitogen.core.Router): def allocate_id(self): return self.id_allocator.allocate() - def context_by_id(self, context_id, via_id=None, create=True): - context = self._context_by_id.get(context_id) - if create and not context: - context = self.context_class(self, context_id) - if via_id is not None: - context.via = self.context_by_id(via_id) - self._context_by_id[context_id] = context - return context - connection_timeout_msg = u"Connection timed out." def _connect(self, klass, name=None, **kwargs):