Merge pull request #220 from Quansight-Labs/getset-state

Add retrieval and setting of state as an opaque object with pickling.

docs/generated/uarray.get_state.rst

@@ -0,0 +1,6 @@
+get\_state
+==========
+
+.. currentmodule:: uarray
+
+.. autofunction:: get_state

docs/generated/uarray.reset_state.rst

@@ -0,0 +1,6 @@
+reset\_state
+============
+
+.. currentmodule:: uarray
+
+.. autofunction:: reset_state

docs/generated/uarray.rst

@@ -20,6 +20,9 @@ uarray
       clear_backends
       skip_backend
       wrap_single_convertor
+      get_state
+      set_state
+      reset_state
 
 
 

docs/generated/uarray.set_state.rst

@@ -0,0 +1,6 @@
+set\_state
+==========
+
+.. currentmodule:: uarray
+
+.. autofunction:: set_state

uarray/_backend.py

@@ -1,10 +1,12 @@
 import typing
+import types
 import inspect
 import functools
 from . import _uarray  # type: ignore
 import copyreg  # type: ignore
 import atexit
 import pickle
+import contextlib
 
 ArgumentExtractorType = typing.Callable[..., typing.Tuple["Dispatchable", ...]]
 ArgumentReplacerType = typing.Callable[
@@ -16,6 +18,7 @@
     _Function,
     _SkipBackendContext,
     _SetBackendContext,
+    _BackendState,
 )
 
 __all__ = [
@@ -32,15 +35,28 @@
     "wrap_single_convertor",
     "all_of_type",
     "mark_as",
+    "set_state",
+    "get_state",
+    "reset_state",
+    "_BackendState",
+    "_SkipBackendContext",
+    "_SetBackendContext",
 ]
 
 
-def unpickle_function(mod_name, qname):
+def unpickle_function(mod_name, qname, self_):
     import importlib
 
     try:
         module = importlib.import_module(mod_name)
-        func = getattr(module, qname)
+        qname = qname.split(".")
+        func = module
+        for q in qname:
+            func = getattr(func, q)
+
+        if self_ is not None:
+            func = types.MethodType(func, self_)
+
         return func
     except (ImportError, AttributeError) as e:
         from pickle import UnpicklingError
@@ -51,9 +67,10 @@ def unpickle_function(mod_name, qname):
 def pickle_function(func):
     mod_name = getattr(func, "__module__", None)
     qname = getattr(func, "__qualname__", None)
+    self_ = getattr(func, "__self__", None)
 
     try:
-        test = unpickle_function(mod_name, qname)
+        test = unpickle_function(mod_name, qname, self_)
     except pickle.UnpicklingError:
         test = None
 
@@ -62,13 +79,76 @@ def pickle_function(func):
             "Can't pickle {}: it's not the same object as {}".format(func, test)
         )
 
-    return unpickle_function, (mod_name, qname)
+    return unpickle_function, (mod_name, qname, self_)
+
+
+def pickle_state(state):
+    return _uarray._BackendState._unpickle, state._pickle()
+
+
+def pickle_set_backend_context(ctx):
+    return _SetBackendContext, ctx._pickle()
+
+
+def pickle_skip_backend_context(ctx):
+    return _SkipBackendContext, ctx._pickle()
 
 
 copyreg.pickle(_Function, pickle_function)
+copyreg.pickle(_uarray._BackendState, pickle_state)
+copyreg.pickle(_SetBackendContext, pickle_set_backend_context)
+copyreg.pickle(_SkipBackendContext, pickle_skip_backend_context)
 atexit.register(_uarray.clear_all_globals)
 
 
+def get_state():
+    """
+    Returns an opaque object containing the current state of all the backends.
+
+    Can be used for synchronization between threads/processes.
+
+    See Also
+    --------
+    set_state
+        Sets the state returned by this function.
+    """
+    return _uarray.get_state()
+
+
+@contextlib.contextmanager
+def reset_state():
+    """
+    Returns a context manager that resets all state once exited.
+
+    See Also
+    --------
+    set_state
+        Context manager that sets the backend state.
+    get_state
+        Gets a state to be set by this context manager.
+    """
+    with set_state(get_state()):
+        yield
+
+
+@contextlib.contextmanager
+def set_state(state):
+    """
+    A context manager that sets the state of the backends to one returned by :obj:`get_state`.
+
+    See Also
+    --------
+    get_state
+        Gets a state to be set by this context manager.
+    """
+    old_state = get_state()
+    _uarray.set_state(state)
+    try:
+        yield
+    finally:
+        _uarray.set_state(old_state, True)
+
+
 def create_multimethod(*args, **kwargs):
     """
     Creates a decorator for generating multimethods.