avoid using contextmanager in Cython
@@ -345,10 +345,7 @@ cdef class _CallbackManager:
self.cb_store_aux_arr = cb_store_aux_arr
-@contextlib.contextmanager
-def set_cufft_callbacks(
- cb_load='', cb_store='', *,
- cb_load_aux_arr=None, cb_store_aux_arr=None):
+cdef class set_cufft_callbacks:
"""A context manager for setting up load and/or store callbacks.
Args:
@@ -420,13 +417,30 @@ def set_cufft_callbacks(
compiled module from cache.
"""
- cdef _ThreadLocal tls = _ThreadLocal.get()
- cdef _CallbackManager mgr, mgr_prev
- cdef tuple key = (cb_load, cb_store)
+ # this class should have been a simple function decorated by @contextlib.
+ # contextmanager that yields...
- try:
- mgr_prev = tls._current_cufft_callback
- mgr = _callback_mgr.get(key)
+ cdef:
+ str cb_load
+ str cb_store
+ ndarray cb_load_aux_arr
+ ndarray cb_store_aux_arr
+ _CallbackManager mgr
+ _CallbackManager mgr_prev
+
+ def __init__(self,
+ str cb_load='',
+ str cb_store='',
+ *,
+ ndarray cb_load_aux_arr=None,
+ ndarray cb_store_aux_arr=None):
+ self.cb_load = cb_load
+ self.cb_store = cb_store
+ self.cb_load_aux_arr = cb_load_aux_arr
+ self.cb_store_aux_arr = cb_store_aux_arr
+
+ cdef tuple key = (cb_load, cb_store)
+ cdef _CallbackManager mgr = _callback_mgr.get(key)
if mgr is None:
mgr = _CallbackManager(
cb_load=cb_load,
@@ -438,9 +452,16 @@ def set_cufft_callbacks(
mgr.set_caller_infos(
cb_load_aux_arr=cb_load_aux_arr,
cb_store_aux_arr=cb_store_aux_arr)
- tls._current_cufft_callback = mgr
- yield mgr
- finally:
- tls._current_cufft_callback = mgr_prev
+ self.mgr = mgr
+
+ def __enter__(self):
+ cdef _ThreadLocal tls = _ThreadLocal.get()
+ self.mgr_prev = tls._current_cufft_callback
+ tls._current_cufft_callback = self.mgr
+ return self.mgr
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ cdef _ThreadLocal tls = _ThreadLocal.get()
+ tls._current_cufft_callback = self.mgr_prev
# do not remove mgr from _callback_mgr, as one might still wanna use
# plans generated by it in the same Python session
