WIP: Namespaced AppContext

petsctools/__init__.py

@@ -15,6 +15,10 @@
 # is not available then attempting to access these attributes will raise an
 # informative error.
 if PETSC4PY_INSTALLED:
+    from .appctx import (  # noqa: F401
+        AppContext,
+        AppContextManager,
+    )
     from .citation import (  # noqa: F401
         add_citation,
         cite,
@@ -65,6 +69,8 @@ def __getattr__(name):
             "set_default_parameter",
             "DefaultOptionSet",
             "PCBase",
+            "AppContext",
+            "AppContextManager",
         }
         if name in petsc4py_attrs:
             raise ImportError(

petsctools/appctx.py

@@ -0,0 +1,197 @@
+from typing import Any
+import itertools
+from functools import cached_property
+from contextlib import contextmanager
+from petsctools.exceptions import PetscToolsAppctxException
+
+_global_appctx_data = {}
+"""The global storage for user data with arbitrary python types."""
+
+
+class AppContextKey(str):
+    """A custom key type for AppContext."""
+
+    _count = itertools.count()
+
+    @classmethod
+    def _generate_key(cls):
+        return f"petsctools_appctx_key_{next(cls._count)}"
+
+
+class AppContext:
+    def __init__(self, prefix: str | None = None):
+        from petsctools.options import _validate_prefix
+
+        # possibly append underscore or cast to str
+        self._prefix = _validate_prefix(prefix or "")
+
+    @property
+    def prefix(self) -> str:
+        return self._prefix
+
+    @cached_property
+    def options_object(self):
+        """A PETSc.Options instance."""
+        from petsc4py import PETSc
+
+        return PETSc.Options()
+
+    def _key_from_option(self, option: str) -> AppContextKey:
+        """
+        Return the internal key for the PETSc option `option`.
+
+        Parameters
+        ----------
+        option
+            The PETSc option.
+
+        Returns
+        -------
+        key
+            An internal key corresponding to ``option``.
+        """
+        return AppContextKey(
+            self.options_object.getString(self.prefix + option)
+        )
+
+    def __getitem__(self, option: str | AppContextKey, /) -> Any:
+        """
+        Return the value with the key saved in ``PETSc.Options()[option]``.
+
+        Parameters
+        ----------
+        option :
+            The PETSc option or key.
+
+        Returns
+        -------
+        val :
+            The value for the key `option`.
+
+        Raises
+        ------
+        PetscToolsAppctxException
+            If the AppContext does contain a value for `option`.
+        """
+        try:
+            return _global_appctx_data[self._key_from_option(option)]
+        except KeyError:
+            raise PetscToolsAppctxException(
+                f"AppContext does not have an entry for {option}"
+            )
+
+    def __setitem__(self, option: str, value: Any, /):
+        key = AppContextKey._generate_key()
+        self.options_object[self.prefix + option] = key
+        _global_appctx_data[key] = value
+
+    def get(
+        self, option: str | AppContextKey, default: Any | None = None
+    ) -> Any:
+        """
+        Return the value with the key saved in ``PETSc.Options()[option]``,
+        or if it does not exist return default.
+
+        Parameters
+        ----------
+        option :
+            The PETSc option or key.
+        default :
+            The value to return if ``option`` is not in the ``AppContext``
+
+        Returns
+        -------
+        val :
+            The value for the key ``option``, or ``default``.
+        """
+        try:
+            return self[option]
+        except PetscToolsAppctxException:
+            return default
+
+
+class AppContextManager:
+    """
+    Class for passing non-primitive types to PETSc python contexts.
+
+    The PETSc.Options dictionary can only contain primitive types (str,
+    int, float, bool) as values. The AppContext allows other types to be
+    passed into PETSc solvers while still making use of the namespacing
+    provided by options prefixing.
+
+    A typical usage is shown below. In this example we have a python PC
+    type `MyCustomPC` which requires additional data in the form of a
+    `MyCustomData` instance.
+    We can add the data to the AppContext with the `appctx.add` method,
+    but we need to tell `MyCustomPC` how to retrieve that data. The
+    `add` method returns a key which is a valid PETSc.Options entry,
+    i.e. a primitive type instance. This key is passed via PETSc.Options
+    with the 'custompc_somedata' prefix.
+
+    NB: The user should never handle this key directly, it should only
+        ever be placed directly into the options dictionary.
+
+    The data can be retrieved by giving the AppContext the (fully
+    prefixed) option for the key, in which case the AppContext will
+    internally fetch the key from the PETSc.Options and return the data.
+
+    .. code-block:: python3
+
+        appctx = AppContext()
+        some_data = MyCustomData(5)
+
+        opts = OptionsManager(
+            parameters={
+                'pc_type': 'python',
+                'pc_python_type': 'MyCustomPC',
+                'custompc_somedata': appctx.add(some_data)},
+            options_prefix='solver')
+
+        with opts.inserted_options():
+            default = MyCustomData(10)
+            data = appctx.get('solver_custompc_somedata', default)
+    """
+
+    def __init__(self):
+        self._data = {}
+
+    def add(self, val: Any) -> AppContextKey:
+        """
+        Add a value to the application context and
+        return the autogenerated key for that value.
+
+        The autogenerated key should be used as the value for the
+        corresponding entry in the solver_parameters dictionary.
+
+        Parameters
+        ----------
+        val
+            The value to add to the AppContext.
+
+        Returns
+        -------
+        key
+            The key to put into the PETSc Options dictionary.
+        """
+        key = AppContextKey._generate_key()
+        self._data[key] = val
+        return key
+
+    @contextmanager
+    def inserted_appctx(self):
+        # We don't overwrite existing entries in the global data,
+        # so we need to keep track of what we do actually put in
+        # so we don't accidentally remove something we shouldn't.
+        to_delete = set()
+        try:
+            for k, v in self._data.items():
+                if k not in _global_appctx_data:
+                    _global_appctx_data[k] = v
+                    to_delete.add(k)
+            yield
+        finally:
+            for k in self._data:
+                if k in to_delete:
+                    del _global_appctx_data[k]
+                    to_delete.remove(k)
+            assert len(to_delete) == 0

petsctools/exceptions.py

@@ -6,5 +6,9 @@ class PetscToolsNotInitialisedException(PetscToolsException):
     """Exception raised when petsctools should have been initialised."""
 
 
+class PetscToolsAppctxException(PetscToolsException):
+    """Exception raised when the Appctx is missing an entry."""
+
+
 class PetscToolsWarning(UserWarning):
     """Generic base class for petsctools warnings."""