test_generator.py 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. import unittest
  2. from greenlet import greenlet
  3. class genlet(greenlet):
  4. def __init__(self, *args, **kwds):
  5. self.args = args
  6. self.kwds = kwds
  7. def run(self):
  8. fn, = self.fn
  9. fn(*self.args, **self.kwds)
  10. def __iter__(self):
  11. return self
  12. def __next__(self):
  13. self.parent = greenlet.getcurrent()
  14. result = self.switch()
  15. if self:
  16. return result
  17. else:
  18. raise StopIteration
  19. # Hack: Python < 2.6 compatibility
  20. next = __next__
  21. def Yield(value):
  22. g = greenlet.getcurrent()
  23. while not isinstance(g, genlet):
  24. if g is None:
  25. raise RuntimeError('yield outside a genlet')
  26. g = g.parent
  27. g.parent.switch(value)
  28. def generator(func):
  29. class generator(genlet):
  30. fn = (func,)
  31. return generator
  32. # ____________________________________________________________
  33. class GeneratorTests(unittest.TestCase):
  34. def test_generator(self):
  35. seen = []
  36. def g(n):
  37. for i in range(n):
  38. seen.append(i)
  39. Yield(i)
  40. g = generator(g)
  41. for k in range(3):
  42. for j in g(5):
  43. seen.append(j)
  44. self.assertEqual(seen, 3 * [0, 0, 1, 1, 2, 2, 3, 3, 4, 4])