python doctest & coroutine

Question

I am having a bit of trouble with doctest and coroutines...

def coroutine(func):
    def start(*args, **kwargs):
        cr=func(*args, **kwargs)
        cr.next()
        return cr
    start.__name__=func.__name__
    return start

@coroutine
def leader_tracking():
    """
    Tracks 'leader' status - only returns transitions

    >>> lt=leader_tracking()
    >>> print lt.send(False)

    >>> print lt.send(False)

    """
    last_status=False
    result=("nop", None)

    while True:
        status=(yield result)

        if status!=last_status:
            direction="up" if last_status==False else "down"
            last_status=status
            result=("tr", direction)
        else:
            result=("nop", None)

If I use the usual doctest scaffolding:

if __name__=="__main__":
    import doctest
    doctest.testmod()

doctest doesn't show anything whilst if I use a more brute force method:

lt=leader_tracking()
print lt.send(True)
print lt.send(False)
print lt.send(False)
print lt.send(True)
print lt.send(True)
print lt.send(True)
print lt.send(True)
print lt.send(False)

I can see the expected result:

('tr', 'up')
('tr', 'down')
('nop', None)
('tr', 'up')
('nop', None)
('nop', None)
('nop', None)
('tr', 'down')

What am I doing wrong with doctest ?

Answer 1

The doctest module looks at docstrings which are stored in the __doc__ attribute of a function. Your coroutine decorator only copies __name__, so the docstring gets lost and doctest can't find it. You could manually assign the __doc__ attribute, but a better way is to use the functools module:

import functools

def coroutine(func):
    @functools.wraps(func)
    def start(*args, **kwargs):
        cr = func(*args, **kwargs)
        cr.next()
        return cr
    return start

The functools.wraps decorator copies all the various metadata attributes of func into the wrapping function start, including the docstring, so that doctest will work as expected.

Also note that your doctest comment should include the expected output in order to work properly:

@coroutine
def leader_tracking():
    """
    Tracks 'leader' status - only returns transitions

    >>> lt=leader_tracking()
    >>> print lt.send(True)
    ('tr', 'up')
    >>> print lt.send(False)
    ('tr', 'down')
    >>> print lt.send(False)
    ('nop', None)
    """
    last_status = False
    result = ("nop", None)

    while True:
        status = yield result

        if status != last_status:
            direction = "up" if last_status == False else "down"
            last_status = status
            result = ("tr", direction)
        else:
            result = ("nop", None)