Keep allocate_id() in the enhanced router class.

mitogen/core.py

@@ -868,11 +868,6 @@ class Router(object):
 
         return handle
 
-    def allocate_id(self):
-        master = Context(self, 0)
-        reply = master.send_await(Message(dst_id=0, handle=ALLOCATE_ID))
-        return reply.unpickle()
-
     def on_shutdown(self, broker):
         """Called during :py:meth:`Broker.shutdown`, informs callbacks
         registered with :py:meth:`add_handle_cb` the connection is dead."""

mitogen/master.py

@@ -739,6 +739,7 @@ METHOD_NAMES = {
 def upgrade_router(econtext):
     if not isinstance(econtext.router, Router):  # TODO
         econtext.router.__class__ = Router  # TODO
+        econtext.router.id_allocator = ChildIdAllocator(econtext.router)
         LOG.debug('_proxy_connect(): constructing ModuleForwarder')
         ModuleForwarder(econtext.router, econtext.parent, econtext.importer)
 
@@ -785,6 +786,18 @@ class IdAllocator(object):
         )
 
 
+class ChildIdAllocator(object):
+    def __init__(self, router):
+        self.router = router
+
+    def allocate(self):
+        master = Context(self.router, 0)
+        reply = master.send_await(
+            mitogen.core.Message(dst_id=0, handle=mitogen.core.ALLOCATE_ID)
+        )
+        return reply.unpickle()
+
+
 class Router(mitogen.core.Router):
     debug = False
 

tests/data/id_allocation.py

 
 import mitogen.core
+import mitogen.master
 
 
-@mitogen.core.takes_router
-def allocate_an_id(router):
-    return router.allocate_id()
+@mitogen.core.takes_econtext
+def allocate_an_id(econtext):
+    mitogen.master.upgrade_router(econtext)
+    return econtext.router.allocate_id()