Search code examples
pythonpython-3.xgeneratorweak-references

How can I get a weak reference to the send method of a generator?


The weakref documentation doesn't seem to provide a method for creating a weak reference to the send method of a generator:

import weakref

def gen(): yield

g=gen()
w_send=weakref.ref(g.send)
w_send() # <- this is None; the g.send object is ephemeral

I didn't think it would work but I did try the weakref.WeakMethod just in case as well:

>>> w_send=weakref.WeakMethod(g.send)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\Users\ricky\AppData\Local\Programs\Python\Python37\lib\weakref.py", line 50, in __new__
    .format(type(meth))) from None
TypeError: argument should be a bound method, not <class 'builtin_function_or_method'>

How can this be done without wrapping the generator in a custom class? Like this:

import weakref

def gen(): yield

class MyGenerator:
    def __init__(self):
        self._generator = gen()
    def send(self, arg):
        return self._generator.send(arg)

g = MyGenerator()
ref = weakref.WeakMethod(g.send)

I don't want to do this. Is there a better way?


The reason I want to do this is I am working on an idea for a simple messaging protocol for an app I might build. The messaging looks something like this:

# messaging module

from typing import Generator
from functools import wraps
from collections import NamedTuple
import weakref

class Message(NamedTuple):
    channel: int
    content: str

_REGISTRY = {}

def _do_register(channel, route):
    # handle bound methods
    if hasattr(route,"__self__") and hasattr(route,"__func__"):
        route_ref = weakref.WeakMethod(route)
    # handle generators
    elif isinstance(route, Generator):
        route_ref = weakref.ref(route.send) # get weak ref to route.send here
    # all other callables
    else:
        route_ref = weakref.ref(route)
    try:
        _REGISTRY[channel].add(route_ref)
    except KeyError:
        _REGISTRY[channel] = {route_ref}

def register(obj=None, *, channel, route=None):
    """Decorator for registering callable objects for messaging."""
    if obj is None:
        def wrapper(callable):
            @wraps(callable)
            def wrapped(*args, **kwargs):
                nonlocal route
                obj_ = callable(*args, **kwargs)
                route_ = obj_ if route is None else route
                _do_register(channel, route_)
                return obj_
            return wrapped
        return wrapper
    else:
        if route is None:
            route = obj
        _do_register(channel, route)

def manager():
    msg_obj = None
    while True:
        msg_obj = yield _broadcast(msg_obj)

def _broadcast(msg_obj):
    count = 0
    if msg_obj:
        for route_ref in _REGISTRY[msg_obj.channel]:
            route = route_ref()
            if route is not None:
                count += 1
                route(msg_obj)
    return count

...used like this:

@register(channel=1)
def listening_gen(name):
    while True:
        msg = yield
        print(f"{name} received message {msg.content} on channel {msg.channel}")


a = listening_gen("a")
b = listening_gen("b")
next(a)
next(b)
register(a, channel=2)
register(b, channel=3)

msg1 = Message(channel=1, content="foo")
msg2 = Message(channel=2, content="bar")
msg3 = Message(channel=3, content="baz")

m = manager()
next(m)
m.send(msg1)
m.send(msg2)
m.send(msg3)

a hears messages on channels 1 and 2, b hears messages on channels 1 and 3.


Solution

  • From the docs:

    Not all objects can be weakly referenced; those objects which can include class instances, functions written in Python (but not in C), instance methods, sets, frozensets, some file objects, generators, type objects, sockets, arrays, deques, regular expression pattern objects, and code objects.

    Since generators are a builtin type written in C, you cannot create a weak reference to a generator's send method. The workaround, as you've already discovered, is to wrap the generator in a python class.