@@ -37,9 +37,11 @@ def is_exit_exception(e):
37
37
38
38
39
39
class _AsyncIoGreenlet (greenlet .greenlet ):
40
+
41
+ __sqlalchemy_greenlet_provider__ = True
42
+
40
43
def __init__ (self , fn , driver ):
41
44
greenlet .greenlet .__init__ (self , fn , driver )
42
- self .driver = driver
43
45
if _has_gr_context :
44
46
self .gr_context = driver .gr_context
45
47
@@ -55,7 +57,7 @@ def await_only(awaitable: Coroutine) -> Any:
55
57
"""
56
58
# this is called in the context greenlet while running fn
57
59
current = greenlet .getcurrent ()
58
- if not isinstance (current , _AsyncIoGreenlet ):
60
+ if not getattr (current , "__sqlalchemy_greenlet_provider__" , False ):
59
61
raise exc .MissingGreenlet (
60
62
"greenlet_spawn has not been called; can't call await_only() "
61
63
"here. Was IO attempted in an unexpected place?"
@@ -65,7 +67,7 @@ def await_only(awaitable: Coroutine) -> Any:
65
67
# a coroutine to run. Once the awaitable is done, the driver greenlet
66
68
# switches back to this greenlet with the result of awaitable that is
67
69
# then returned to the caller (or raised as error)
68
- return current .driver .switch (awaitable )
70
+ return current .parent .switch (awaitable )
69
71
70
72
71
73
def await_fallback (awaitable : Coroutine ) -> Any :
@@ -79,7 +81,7 @@ def await_fallback(awaitable: Coroutine) -> Any:
79
81
"""
80
82
# this is called in the context greenlet while running fn
81
83
current = greenlet .getcurrent ()
82
- if not isinstance (current , _AsyncIoGreenlet ):
84
+ if not getattr (current , "__sqlalchemy_greenlet_provider__" , False ):
83
85
loop = get_event_loop ()
84
86
if loop .is_running ():
85
87
raise exc .MissingGreenlet (
@@ -89,7 +91,7 @@ def await_fallback(awaitable: Coroutine) -> Any:
89
91
)
90
92
return loop .run_until_complete (awaitable )
91
93
92
- return current .driver .switch (awaitable )
94
+ return current .parent .switch (awaitable )
93
95
94
96
95
97
async def greenlet_spawn (
@@ -111,24 +113,21 @@ async def greenlet_spawn(
111
113
# coroutine to wait. If the context is dead the function has
112
114
8000
# returned, and its result can be returned.
113
115
switch_occurred = False
114
- try :
115
- result = context .switch (* args , ** kwargs )
116
- while not context .dead :
117
- switch_occurred = True
118
- try :
119
- # wait for a coroutine from await_only and then return its
120
- # result back to it.
121
- value = await result
122
- except BaseException :
123
- # this allows an exception to be raised within
124
- # the moderated greenlet so that it can continue
125
- # its expected flow.
126
- result = context .throw (* sys .exc_info ())
127
- else :
128
- result = context .switch (value )
129
- finally :
130
- # clean up to avoid cycle resolution by gc
131
- del context .driver
116
+ result = context .switch (* args , ** kwargs )
117
+ while not context .dead :
118
+ switch_occurred = True
119
+ try :
120
+ # wait for a coroutine from await_only and then return its
121
+ # result back to it.
122
+ value = await result
123
+ except BaseException :
124
+ # this allows an exception to be raised within
125
+ # the moderated greenlet so that it can continue
126
+ # its expected flow.
127
+ result = context .throw (* sys .exc_info ())
128
+ else :
129
+ result = context .switch (value )
130
+
132
131
if _require_await and not switch_occurred :
133
132
raise exc .AwaitRequired (
134
133
"The current operation required an async execution but none was "
@@ -175,7 +174,9 @@ def _util_async_run(fn, *args, **kwargs):
175
174
return loop .run_until_complete (greenlet_spawn (fn , * args , ** kwargs ))
176
175
else :
177
176
# allow for a wrapped test function to call another
178
- assert isinstance (greenlet .getcurrent (), _AsyncIoGreenlet )
177
+ assert getattr (
178
+ greenlet .getcurrent (), "__sqlalchemy_greenlet_provider__" , False
179
+ )
179
180
return fn (* args , ** kwargs )
180
181
181
182
0 commit comments