test_setcomps.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160
  1. import doctest
  2. import unittest
  3. doctests = """
  4. ########### Tests mostly copied from test_listcomps.py ############
  5. Test simple loop with conditional
  6. >>> sum({i*i for i in range(100) if i&1 == 1})
  7. 166650
  8. Test simple case
  9. >>> {2*y + x + 1 for x in (0,) for y in (1,)}
  10. {3}
  11. Test simple nesting
  12. >>> list(sorted({(i,j) for i in range(3) for j in range(4)}))
  13. [(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3), (2, 0), (2, 1), (2, 2), (2, 3)]
  14. Test nesting with the inner expression dependent on the outer
  15. >>> list(sorted({(i,j) for i in range(4) for j in range(i)}))
  16. [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)]
  17. Test the idiom for temporary variable assignment in comprehensions.
  18. >>> sorted({j*j for i in range(4) for j in [i+1]})
  19. [1, 4, 9, 16]
  20. >>> sorted({j*k for i in range(4) for j in [i+1] for k in [j+1]})
  21. [2, 6, 12, 20]
  22. >>> sorted({j*k for i in range(4) for j, k in [(i+1, i+2)]})
  23. [2, 6, 12, 20]
  24. Not assignment
  25. >>> sorted({i*i for i in [*range(4)]})
  26. [0, 1, 4, 9]
  27. >>> sorted({i*i for i in (*range(4),)})
  28. [0, 1, 4, 9]
  29. Make sure the induction variable is not exposed
  30. >>> i = 20
  31. >>> sum({i*i for i in range(100)})
  32. 328350
  33. >>> i
  34. 20
  35. Verify that syntax error's are raised for setcomps used as lvalues
  36. >>> {y for y in (1,2)} = 10 # doctest: +IGNORE_EXCEPTION_DETAIL
  37. Traceback (most recent call last):
  38. ...
  39. SyntaxError: ...
  40. >>> {y for y in (1,2)} += 10 # doctest: +IGNORE_EXCEPTION_DETAIL
  41. Traceback (most recent call last):
  42. ...
  43. SyntaxError: ...
  44. Make a nested set comprehension that acts like set(range())
  45. >>> def srange(n):
  46. ... return {i for i in range(n)}
  47. >>> list(sorted(srange(10)))
  48. [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
  49. Same again, only as a lambda expression instead of a function definition
  50. >>> lrange = lambda n: {i for i in range(n)}
  51. >>> list(sorted(lrange(10)))
  52. [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
  53. Generators can call other generators:
  54. >>> def grange(n):
  55. ... for x in {i for i in range(n)}:
  56. ... yield x
  57. >>> list(sorted(grange(5)))
  58. [0, 1, 2, 3, 4]
  59. Make sure that None is a valid return value
  60. >>> {None for i in range(10)}
  61. {None}
  62. ########### Tests for various scoping corner cases ############
  63. Return lambdas that use the iteration variable as a default argument
  64. >>> items = {(lambda i=i: i) for i in range(5)}
  65. >>> {x() for x in items} == set(range(5))
  66. True
  67. Same again, only this time as a closure variable
  68. >>> items = {(lambda: i) for i in range(5)}
  69. >>> {x() for x in items}
  70. {4}
  71. Another way to test that the iteration variable is local to the list comp
  72. >>> items = {(lambda: i) for i in range(5)}
  73. >>> i = 20
  74. >>> {x() for x in items}
  75. {4}
  76. And confirm that a closure can jump over the list comp scope
  77. >>> items = {(lambda: y) for i in range(5)}
  78. >>> y = 2
  79. >>> {x() for x in items}
  80. {2}
  81. We also repeat each of the above scoping tests inside a function
  82. >>> def test_func():
  83. ... items = {(lambda i=i: i) for i in range(5)}
  84. ... return {x() for x in items}
  85. >>> test_func() == set(range(5))
  86. True
  87. >>> def test_func():
  88. ... items = {(lambda: i) for i in range(5)}
  89. ... return {x() for x in items}
  90. >>> test_func()
  91. {4}
  92. >>> def test_func():
  93. ... items = {(lambda: i) for i in range(5)}
  94. ... i = 20
  95. ... return {x() for x in items}
  96. >>> test_func()
  97. {4}
  98. >>> def test_func():
  99. ... items = {(lambda: y) for i in range(5)}
  100. ... y = 2
  101. ... return {x() for x in items}
  102. >>> test_func()
  103. {2}
  104. """
  105. __test__ = {'doctests' : doctests}
  106. def load_tests(loader, tests, pattern):
  107. tests.addTest(doctest.DocTestSuite())
  108. return tests
  109. if __name__ == "__main__":
  110. unittest.main()