Generalize delegation to allow prefixes

pyqg/__init__.py

@@ -5,6 +5,7 @@
 from .layered_model import LayeredModel
 from .particles import LagrangianParticleArray2D, GriddedLagrangianParticleArray2D
 from .parameterizations import *
+from .delegate import delegate
 
 try:
     from importlib.metadata import version, PackageNotFoundError

pyqg/vendor/delegate/__init__.py → pyqg/delegate.py

@@ -1,18 +1,21 @@
-# Taken from https://github.com/dscottboggs/python-delegate
-
-def bind(instance, func, name):
-    """Turn a function into a bound method on instance"""
-    setattr(instance, name, func.__get__(instance, instance.__class__))
+# Adapted from https://github.com/dscottboggs/python-delegate
 
 def delegate(*args, **named_args):
-    dest = named_args.get('to')
-    if dest is None:
+    _dest = named_args.get('to')
+    if _dest is None:
         raise ValueError(
             "the named argument 'to' is required on the delegate function")
+
+    _prefix = named_args.get('prefix', '')
+
+    print(_dest, _prefix)
+
     def wraps(cls, *wrapped_args, **wrapped_opts):
         """Wrap the target class up in something that modifies."""
         class Wrapped(cls):
 
+            _delegates = [(a, _dest, _prefix) for a in args] + getattr(cls, '_delegates', [])
+
             def __getattr__(self, name):
                 """
                 Return the selected name from the destination if the name is one
@@ -21,24 +24,34 @@ def __getattr__(self, name):
                 error when `name` is not one of the selected args to be
                 delegated.
                 """
-                if name in args: return getattr(self.__dict__[dest], name)
+                for attr, dest, prefix in self._delegates:
+                    if name == prefix + attr:
+                        return getattr(self.__dict__[dest], name[len(prefix):])
+
                 raise AttributeError(f"'{self.__class__.__name__}' has no attribute '{name}'")
 
             def __setattr__(self, name, value):
                 """
                 If this name is one of those selected, set it on the destination
                 property. Otherwise, set it on self.
                 """
-                if name in args: setattr(getattr(self, dest), name, value)
-                else: self.__dict__[name] = value
+                for attr, dest, prefix in self._delegates:
+                    if name == prefix + attr:
+                        setattr(getattr(self, dest), name[len(prefix):], value)
+                        return
+                self.__dict__[name] = value
 
             def __delattr__(self, name):
                 """Delete name from `dest` or `self`"""
-                if name in args: delattr(getattr(self, dest), name)
-                else: del self.__dict__[name]
+                for attr, dest, prefix in self._delegates:
+                    if name == prefix + attr:
+                        delattr(getattr(self, dest), name[len(prefix):])
+                        return
+
+                del self.__dict__[name]
 
         Wrapped.__doc__ = cls.__doc__ or \
-            f"{cls.__class__} wrapped to delegate {args} to its {dest} property"
+            f"{cls.__class__} wrapped to delegate {args} to its {_dest} property"
         Wrapped.__repr__ = cls.__repr__
         Wrapped.__str__ = cls.__str__
         Wrapped.__name__ = cls.__name__

pyqg/model.py

@@ -6,7 +6,7 @@
 
 from .errors import DiagnosticNotFilledError
 from .parameterizations import Parameterization
-from .vendor.delegate import delegate
+from .delegate import delegate
 from . import kernels
 
 try:
@@ -21,14 +21,20 @@
 except ImportError:
     pass
 
-kernel_attrs = [
+public_kernel_attrs = [
     'nx', 'ny', 'nz', 'nl', 'nk', 'a', 'kk', 'll', 'q', 'qh', 'ph', 'u', 'v',
     'Ubg', 'Qy', 'ufull', 'vfull', 'rek', 't', 'tc', 'dt',
     'uv_parameterization', 'q_parameterization', 'fft', 'ifft', 'filtr',
     'ablevel', 'dqhdt', 'dqhdt_p', 'dqhdt_pp', 'dqh', 'duh', 'dvh',
 ]
 
-@delegate(*kernel_attrs, to='kernel')
+private_kernel_attrs = [
+    'invert', 'do_advection', 'do_friction', 'do_q_subgrid_parameterization',
+    'do_uv_subgrid_parameterization', 'forward_timestep',
+]
+
+@delegate(*public_kernel_attrs, to='kernel')
+@delegate(*private_kernel_attrs, to='kernel', prefix='_')
 class Model:
     """A generic pseudo-spectral inversion model.
 
@@ -437,24 +443,6 @@ def _step_forward(self):
         # the basic steps are
         self._print_status()
 
-    def _invert(self):
-        self.kernel.invert()
-
-    def _do_advection(self):
-        self.kernel.do_advection()
-
-    def _do_friction(self):
-        self.kernel.do_friction()
-
-    def _do_q_subgrid_parameterization(self):
-        self.kernel.do_q_subgrid_parameterization()
-
-    def _do_uv_subgrid_parameterization(self):
-        self.kernel.do_uv_subgrid_parameterization()
-
-    def _forward_timestep(self):
-        self.kernel.forward_timestep()
-
     def _initialize_time(self):
         """Set up timestep stuff"""
         #self.t=0        # actual time

pyqg/tests/test_delegate.py

@@ -0,0 +1,38 @@
+from pyqg import delegate
+
+def test_delegate():
+    @delegate(*['e', 'f'], to='b')
+    @delegate(*['g'], to='c', prefix='_')
+    @delegate(*['h'], to='c', prefix='__')
+    class A:
+        def __init__(self, b, c):
+            self.b = b
+            self.c = c
+            self.d = 'foo'
+
+    class B:
+        def __init__(self):
+            self.e = 'bar'
+            self.f = 'baz'
+
+    class C:
+        def __init__(self):
+            self.e = 'NOPE'
+            self.g = 'bat'
+            self.h = 'bart'
+
+    delegates = [('e', 'b', ''), ('f', 'b', ''),
+                 ('g', 'c', '_'), ('h', 'c', '__')]
+
+    assert A._delegates == delegates
+
+    a = A(B(), C())
+
+    assert a.d == 'foo'
+    assert a.e == 'bar'
+    assert a.f == 'baz'
+    assert a._g == 'bat'
+    assert a.__h == 'bart'
+
+if __name__ == '__main__':
+    test_delegate()