test_generator_nested.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165
  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. self.child = None
  8. def run(self):
  9. fn, = self.fn
  10. fn(*self.args, **self.kwds)
  11. def __iter__(self):
  12. return self
  13. def set_child(self, child):
  14. self.child = child
  15. def __next__(self):
  16. if self.child:
  17. child = self.child
  18. while child.child:
  19. tmp = child
  20. child = child.child
  21. tmp.child = None
  22. result = child.switch()
  23. else:
  24. self.parent = greenlet.getcurrent()
  25. result = self.switch()
  26. if self:
  27. return result
  28. else:
  29. raise StopIteration
  30. # Hack: Python < 2.6 compatibility
  31. next = __next__
  32. def Yield(value, level=1):
  33. g = greenlet.getcurrent()
  34. while level != 0:
  35. if not isinstance(g, genlet):
  36. raise RuntimeError('yield outside a genlet')
  37. if level > 1:
  38. g.parent.set_child(g)
  39. g = g.parent
  40. level -= 1
  41. g.switch(value)
  42. def Genlet(func):
  43. class Genlet(genlet):
  44. fn = (func,)
  45. return Genlet
  46. # ____________________________________________________________
  47. def g1(n, seen):
  48. for i in range(n):
  49. seen.append(i + 1)
  50. yield i
  51. def g2(n, seen):
  52. for i in range(n):
  53. seen.append(i + 1)
  54. Yield(i)
  55. g2 = Genlet(g2)
  56. def nested(i):
  57. Yield(i)
  58. def g3(n, seen):
  59. for i in range(n):
  60. seen.append(i + 1)
  61. nested(i)
  62. g3 = Genlet(g3)
  63. def a(n):
  64. if n == 0:
  65. return
  66. for ii in ax(n - 1):
  67. Yield(ii)
  68. Yield(n)
  69. ax = Genlet(a)
  70. def perms(l):
  71. if len(l) > 1:
  72. for e in l:
  73. # No syntactical sugar for generator expressions
  74. [Yield([e] + p) for p in perms([x for x in l if x != e])]
  75. else:
  76. Yield(l)
  77. perms = Genlet(perms)
  78. def gr1(n):
  79. for ii in range(1, n):
  80. Yield(ii)
  81. Yield(ii * ii, 2)
  82. gr1 = Genlet(gr1)
  83. def gr2(n, seen):
  84. for ii in gr1(n):
  85. seen.append(ii)
  86. gr2 = Genlet(gr2)
  87. class NestedGeneratorTests(unittest.TestCase):
  88. def test_layered_genlets(self):
  89. seen = []
  90. for ii in gr2(5, seen):
  91. seen.append(ii)
  92. self.assertEqual(seen, [1, 1, 2, 4, 3, 9, 4, 16])
  93. def test_permutations(self):
  94. gen_perms = perms(list(range(4)))
  95. permutations = list(gen_perms)
  96. self.assertEqual(len(permutations), 4 * 3 * 2 * 1)
  97. self.assertTrue([0, 1, 2, 3] in permutations)
  98. self.assertTrue([3, 2, 1, 0] in permutations)
  99. res = []
  100. for ii in zip(perms(list(range(4))), perms(list(range(3)))):
  101. res.append(ii)
  102. self.assertEqual(
  103. res,
  104. [([0, 1, 2, 3], [0, 1, 2]), ([0, 1, 3, 2], [0, 2, 1]),
  105. ([0, 2, 1, 3], [1, 0, 2]), ([0, 2, 3, 1], [1, 2, 0]),
  106. ([0, 3, 1, 2], [2, 0, 1]), ([0, 3, 2, 1], [2, 1, 0])])
  107. # XXX Test to make sure we are working as a generator expression
  108. def test_genlet_simple(self):
  109. for g in [g1, g2, g3]:
  110. seen = []
  111. for k in range(3):
  112. for j in g(5, seen):
  113. seen.append(j)
  114. self.assertEqual(seen, 3 * [1, 0, 2, 1, 3, 2, 4, 3, 5, 4])
  115. def test_genlet_bad(self):
  116. try:
  117. Yield(10)
  118. except RuntimeError:
  119. pass
  120. def test_nested_genlets(self):
  121. seen = []
  122. for ii in ax(5):
  123. seen.append(ii)