codekingpro/portable-devtools

1import gc
2import sys
3import unittest
4from contextvars import Context
5from contextvars import ContextVar
6from contextvars import copy_context
7from functools import partial
8
9from greenlet import getcurrent
10from greenlet import greenlet
11
12from . import PY314
13from . import TestCase
14
15# From the documentation:
16#
17# Important: Context Variables should be created at the top module
18# level and never in closures. Context objects hold strong
19# references to context variables which prevents context variables
20# from being properly garbage collected.
21ID_VAR = ContextVar("id", default=None)
22VAR_VAR = ContextVar("var", default=None)
23ContextVar = None
24
25
26class ContextVarsTests(TestCase):
27    def _new_ctx_run(self, *args, **kwargs):
28        return copy_context().run(*args, **kwargs)
29
30    def _increment(self, greenlet_id, callback, counts, expect):
31        ctx_var = ID_VAR
32        if expect is None:
33            self.assertIsNone(ctx_var.get())
34        else:
35            self.assertEqual(ctx_var.get(), expect)
36        ctx_var.set(greenlet_id)
37        for _ in range(2):
38            counts[ctx_var.get()] += 1
39            callback()
40
41    def _test_context(self, propagate_by):
42        # pylint:disable=too-many-branches
43        ID_VAR.set(0)
44
45        callback = getcurrent().switch
46        counts = dict((i, 0) for i in range(5))
47
48        lets = [
49            greenlet(partial(
50                partial(
51                    copy_context().run,
52                    self._increment
53                ) if propagate_by == "run" else self._increment,
54                greenlet_id=i,
55                callback=callback,
56                counts=counts,
57                expect=(
58                    i - 1 if propagate_by == "share" else
59                    0 if propagate_by in ("set", "run") else None
60                )
61            ))
62            for i in range(1, 5)
63        ]
64
65        for let in lets:
66            if propagate_by == "set":
67                let.gr_context = copy_context()
68            elif propagate_by == "share":
69                let.gr_context = getcurrent().gr_context
70
71        for i in range(2):
72            counts[ID_VAR.get()] += 1
73            for let in lets:
74                let.switch()
75
76        if propagate_by == "run":
77            # Must leave each context.run() in reverse order of entry
78            for let in reversed(lets):
79                let.switch()
80        else:
81            # No context.run(), so fine to exit in any order.
82            for let in lets:
83                let.switch()
84
85        for let in lets:
86            self.assertTrue(let.dead)
87            # When using run(), we leave the run() as the greenlet dies,
88            # and there's no context "underneath". When not using run(),
89            # gr_context still reflects the context the greenlet was
90            # running in.
91            if propagate_by == 'run':
92                self.assertIsNone(let.gr_context)
93            else:
94                self.assertIsNotNone(let.gr_context)
95
96
97        if propagate_by == "share":
98            self.assertEqual(counts, {0: 1, 1: 1, 2: 1, 3: 1, 4: 6})
99        else:
100            self.assertEqual(set(counts.values()), set([2]))
101
102    def test_context_propagated_by_context_run(self):
103        self._new_ctx_run(self._test_context, "run")
104
105    def test_context_propagated_by_setting_attribute(self):
106        self._new_ctx_run(self._test_context, "set")
107
108    def test_context_not_propagated(self):
109        self._new_ctx_run(self._test_context, None)
110
111    def test_context_shared(self):
112        self._new_ctx_run(self._test_context, "share")
113
114    def test_break_ctxvars(self):
115        let1 = greenlet(copy_context().run)
116        let2 = greenlet(copy_context().run)
117        let1.switch(getcurrent().switch)
118        let2.switch(getcurrent().switch)
119        # Since let2 entered the current context and let1 exits its own, the
120        # interpreter emits:
121        # RuntimeError: cannot exit context: thread state references a different context object
122        let1.switch()
123
124    def test_not_broken_if_using_attribute_instead_of_context_run(self):
125        let1 = greenlet(getcurrent().switch)
126        let2 = greenlet(getcurrent().switch)
127        let1.gr_context = copy_context()
128        let2.gr_context = copy_context()
129        let1.switch()
130        let2.switch()
131        let1.switch()
132        let2.switch()
133
134    def test_context_assignment_while_running(self):
135        # pylint:disable=too-many-statements
136        ID_VAR.set(None)
137
138        def target():
139            self.assertIsNone(ID_VAR.get())
140            self.assertIsNone(gr.gr_context)
141
142            # Context is created on first use
143            ID_VAR.set(1)
144            self.assertIsInstance(gr.gr_context, Context)
145            self.assertEqual(ID_VAR.get(), 1)
146            self.assertEqual(gr.gr_context[ID_VAR], 1)
147
148            # Clearing the context makes it get re-created as another
149            # empty context when next used
150            old_context = gr.gr_context
151            gr.gr_context = None  # assign None while running
152            self.assertIsNone(ID_VAR.get())
153            self.assertIsNone(gr.gr_context)
154            ID_VAR.set(2)
155            self.assertIsInstance(gr.gr_context, Context)
156            self.assertEqual(ID_VAR.get(), 2)
157            self.assertEqual(gr.gr_context[ID_VAR], 2)
158
159            new_context = gr.gr_context
160            getcurrent().parent.switch((old_context, new_context))
161            # parent switches us back to old_context
162
163            self.assertEqual(ID_VAR.get(), 1)
164            gr.gr_context = new_context  # assign non-None while running
165            self.assertEqual(ID_VAR.get(), 2)
166
167            getcurrent().parent.switch()
168            # parent switches us back to no context
169            self.assertIsNone(ID_VAR.get())
170            self.assertIsNone(gr.gr_context)
171            gr.gr_context = old_context
172            self.assertEqual(ID_VAR.get(), 1)
173
174            getcurrent().parent.switch()
175            # parent switches us back to no context
176            self.assertIsNone(ID_VAR.get())
177            self.assertIsNone(gr.gr_context)
178
179        gr = greenlet(target)
180
181        with self.assertRaisesRegex(AttributeError, "can't delete context attribute"):
182            del gr.gr_context
183
184        self.assertIsNone(gr.gr_context)
185        old_context, new_context = gr.switch()
186        self.assertIs(new_context, gr.gr_context)
187        self.assertEqual(old_context[ID_VAR], 1)
188        self.assertEqual(new_context[ID_VAR], 2)
189        self.assertEqual(new_context.run(ID_VAR.get), 2)
190        gr.gr_context = old_context  # assign non-None while suspended
191        gr.switch()
192        self.assertIs(gr.gr_context, new_context)
193        gr.gr_context = None  # assign None while suspended
194        gr.switch()
195        self.assertIs(gr.gr_context, old_context)
196        gr.gr_context = None
197        gr.switch()
198        self.assertIsNone(gr.gr_context)
199
200        # Make sure there are no reference leaks
201        gr = None
202        gc.collect()
203        # Python 3.14 elides reference counting operations
204        # in some cases. See https://github.com/python/cpython/pull/130708
205        self.assertEqual(sys.getrefcount(old_context), 2 if not PY314 else 1)
206        self.assertEqual(sys.getrefcount(new_context), 2 if not PY314 else 1)
207
208    def test_context_assignment_different_thread(self):
209        import threading
210        VAR_VAR.set(None)
211        ctx = Context()
212
213        is_running = threading.Event()
214        should_suspend = threading.Event()
215        did_suspend = threading.Event()
216        should_exit = threading.Event()
217        holder = []
218
219        def greenlet_in_thread_fn():
220            VAR_VAR.set(1)
221            is_running.set()
222            should_suspend.wait(10)
223            VAR_VAR.set(2)
224            getcurrent().parent.switch()
225            holder.append(VAR_VAR.get())
226
227        def thread_fn():
228            gr = greenlet(greenlet_in_thread_fn)
229            gr.gr_context = ctx
230            holder.append(gr)
231            gr.switch()
232            did_suspend.set()
233            should_exit.wait(10)
234            gr.switch()
235            del gr
236            greenlet() # trigger cleanup
237
238        thread = threading.Thread(target=thread_fn, daemon=True)
239        thread.start()
240        is_running.wait(10)
241        gr = holder[0]
242
243        # Can't access or modify context if the greenlet is running
244        # in a different thread
245        with self.assertRaisesRegex(ValueError, "running in a different"):
246            getattr(gr, 'gr_context')
247        with self.assertRaisesRegex(ValueError, "running in a different"):
248            gr.gr_context = None
249
250        should_suspend.set()
251        did_suspend.wait(10)
252
253        # OK to access and modify context if greenlet is suspended
254        self.assertIs(gr.gr_context, ctx)
255        self.assertEqual(gr.gr_context[VAR_VAR], 2)
256        gr.gr_context = None
257
258        should_exit.set()
259        thread.join(10)
260
261        self.assertEqual(holder, [gr, None])
262
263        # Context can still be accessed/modified when greenlet is dead:
264        self.assertIsNone(gr.gr_context)
265        gr.gr_context = ctx
266        self.assertIs(gr.gr_context, ctx)
267
268        # Otherwise we leak greenlets on some platforms.
269        # XXX: Should be able to do this automatically
270        del holder[:]
271        gr = None
272        thread = None
273
274    def test_context_assignment_wrong_type(self):
275        g = greenlet()
276        with self.assertRaisesRegex(TypeError,
277                                    "greenlet context must be a contextvars.Context or None"):
278            g.gr_context = self
279
280
281if __name__ == '__main__':
282    unittest.main()
283