1# test that iterating doesn't use the heap
2try:
3    frozenset
4except NameError:
5    print("SKIP")
6    raise SystemExit
7try:
8    import uarray as array
9except ImportError:
10    try:
11        import array
12    except ImportError:
13        print("SKIP")
14        raise SystemExit
15
16try:
17    from micropython import heap_lock, heap_unlock
18except (ImportError, AttributeError):
19    heap_lock = heap_unlock = lambda: 0
20
21
22def do_iter(l):
23    heap_lock()
24    for i in l:
25        print(i)
26    heap_unlock()
27
28
29def gen_func():
30    yield 1
31    yield 2
32
33
34# pre-create collections to iterate over
35ba = bytearray(b"123")
36ar = array.array("H", (123, 456))
37t = (1, 2, 3)
38l = [1, 2]
39d = {1: 2}
40s = set((1,))
41fs = frozenset((1,))
42g1 = (100 + x for x in range(2))
43g2 = gen_func()
44
45# test containment (both success and failure) with the heap locked
46heap_lock()
47print(49 in b"123", 255 in b"123")
48print(1 in t, -1 in t)
49print(1 in l, -1 in l)
50print(1 in d, -1 in d)
51print(1 in s, -1 in s)
52heap_unlock()
53
54# test unpacking with the heap locked
55unp0 = unp1 = unp2 = None  # preallocate slots for globals
56heap_lock()
57unp0, unp1, unp2 = t
58print(unp0, unp1, unp2)
59heap_unlock()
60
61# test certain builtins with the heap locked
62heap_lock()
63print(all(t))
64print(any(t))
65print(min(t))
66print(max(t))
67print(sum(t))
68heap_unlock()
69
70# test iterating over collections with the heap locked
71do_iter(b"123")
72do_iter(ba)
73do_iter(ar)
74do_iter(t)
75do_iter(l)
76do_iter(d)
77do_iter(s)
78do_iter(fs)
79do_iter(g1)
80do_iter(g2)
81