WIP: Namespaced AppContext

petsctools/__init__.py

@@ -1,3 +1,4 @@
+from .appctx import AppContext  # noqa: F401
 from .config import (  # noqa: F401
     MissingPetscException,
     get_config,

petsctools/appctx.py

@@ -0,0 +1,185 @@
+import itertools
+from functools import cached_property
+from petsctools.exceptions import PetscToolsAppctxException
+
+
+class AppContextKey(int):
+    """A custom key type for AppContext."""
+    pass
+
+
+class AppContext:
+    """
+    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.
+
+    The data can be retrieved in two ways.
+    1) 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.
+    2) By manually fetching the AppContext key from the PETSc.Options,
+       then retrieving the data from the `AppContext` using that key.
+
+    .. 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():
+            # 1) Let AppContext fetch key.
+            #    Also shows providing default data.
+            default = MyCustomData(10)
+            data = appctx.get('solver_custompc_somedata', default)
+
+            # 2) Fetch key directly.
+            key = PETSc.Options()['solver_custompc_somedata']
+            data = appctx[key]
+    """
+
+    def __init__(self):
+        self._count = itertools.count()
+        self._data = {}
+
+    def _keygen(self):
+        """
+        Generate a new unique internal key.
+
+        This should not called directly by the user.
+        """
+        return AppContextKey(next(self._count))
+
+    def _to_key(self, option):
+        """
+        Return the internal key for the PETSc option `option`.
+        If `option` is already an AppContextKey, `option` is returned.
+
+        This should not called directly by the user.
+        """
+        if isinstance(option, int):
+            return AppContextKey(option)
+        else:
+            return self.getKey(option)
+
+    @cached_property
+    def _missing_key(self):
+        """
+        Key instance representing a missing AppContext entry.
+
+        PETSc requires the default value for Options.getObj()
+        to be the correct type, so we need a dummy key.
+
+        This should not called directly by the user.
+        """
+        return self._keygen()
+
+    def getKey(self, option):
+        """
+        Return the internal key for the PETSc option `option`.
+
+        Parameters
+        ----------
+        option : str
+            The PETSc option.
+
+        Returns
+        -------
+        key : AppContextKey
+            An internal key corresponding to `option`.
+        """
+        key = self.options_object.getInt(option, self._missing_key)
+        return AppContextKey(key)
+
+    def add(self, val):
+        """
+        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 : Any
+            The value to add to the AppContext.
+
+        Returns
+        -------
+        key : AppContextKey
+            The key to put into the PETSc Options dictionary.
+        """
+        key = self._keygen()
+        self._data[key] = val
+        return key
+
+    def __getitem__(self, option):
+        """
+        Return the value with the key saved in `PETSc.Options()[option]`.
+
+        Parameters
+        ----------
+        option : Union[str, AppContextKey]
+            The PETSc option or key.
+
+        Returns
+        -------
+        val : Any
+            The value for the key `option`.
+
+        Raises
+        ------
+        PetscToolsAppctxException
+            If the AppContext does contain a value for `option`.
+        """
+        try:
+            return self._data[self._to_key(option)]
+        except KeyError:
+            raise PetscToolsAppctxException(
+                f"AppContext does not have an entry for {option}")
+
+    def get(self, option, default=None):
+        """
+        Return the value with the key saved in PETSc.Options()[option],
+        or if it does not exist return default.
+
+        Parameters
+        ----------
+        option : Union[str, AppContextKey]
+            The PETSc option or key.
+        default : Any
+            The value to return if `option` is not in the AppContext
+
+        Returns
+        -------
+        val : Any
+            The value for the key `option`, or `default`.
+        """
+        key = self._to_key(option)
+        if key == self._missing_key:
+            return default
+        return self._data[key]
+
+    @cached_property
+    def options_object(self):
+        """A PETSc.Options instance."""
+        from petsc4py import PETSc
+        return PETSc.Options()

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."""

tests/test_appctx.py

@@ -0,0 +1,41 @@
+import pytest
+import petsctools
+
+
+@pytest.mark.skipnopetsc4py
+def test_appctx():
+    PETSc = petsctools.init()
+
+    appctx = petsctools.AppContext()
+
+    param = 10
+    key = appctx.add(param)
+    options = PETSc.Options()
+    options['solver_param'] = key
+
+    # Can we get the key string back?
+    assert str(appctx.getKey('solver_param')) == options['solver_param']
+
+    # Can we access param via the prefixed option?
+    prm = appctx.get('solver_param')
+    assert prm is param
+
+    prm = appctx['solver_param']
+    assert prm is param
+
+    # Can we access param via the key?
+    prm = appctx.get(key, 20)
+    assert prm is param
+
+    prm = appctx[key]
+    assert prm is param
+
+    # Can we set a default value?
+    default = 20
+    prm = appctx.get('param', default)
+    assert prm is default
+
+    # Will an invalid key raise an error
+    from petsctools.appctx import PetscToolsAppctxException
+    with pytest.raises(PetscToolsAppctxException):
+        appctx['param']