Implement accessor

pydomcfg/__init__.py

@@ -1,11 +1,15 @@
 from pkg_resources import DistributionNotFound, get_distribution
 
 from . import utils
+from .accessor import Accessor
 
 try:
     __version__ = get_distribution("pydomcfg").version
 except DistributionNotFound:
     # package is not installed
     __version__ = "unknown"
 
-__all__ = ("utils",)
+__all__ = (
+    "Accessor",
+    "utils",
+)

pydomcfg/accessor.py

@@ -0,0 +1,46 @@
+from typing import Callable
+
+import xarray as xr
+from xarray import Dataset
+
+from .domzgr.zco import Zco
+
+
+def _property_and_jpk_check(func: Callable) -> Callable:
+    """
+    1. Transform the function to property
+    2. Raise an error if jpk was not set
+    """
+
+    @property  # type: ignore # TODO: "property" used with a non-method.
+    def wrapper(self):
+        if not self.jpk:
+            raise ValueError(
+                f"Set `jpk` before calling `obj.domcfg.{func.__name__}`."
+                " For example: obj.domcfg.jpk = 31"
+            )
+
+        return func(self)
+
+    return wrapper
+
+
+@xr.register_dataset_accessor("domcfg")
+class Accessor:
+    def __init__(self, xarray_obj: Dataset):
+        self._obj = xarray_obj
+        self._jpk: int = 0
+
+    @property
+    def jpk(self):
+        return self._jpk
+
+    @jpk.setter
+    def jpk(self, value: int):
+        if value <= 0:
+            raise ValueError("`jpk` MUST be > 0")
+        self._jpk = value
+
+    @_property_and_jpk_check
+    def zco(self) -> Callable:
+        return Zco(self._obj, self._jpk).__call__

pydomcfg/tests/test_zco.py

@@ -6,7 +6,7 @@
 import pytest
 import xarray as xr
 
-from pydomcfg.domzgr.zco import Zco
+import pydomcfg  # noqa
 
 from .bathymetry import Bathymetry
 from .data import ORCA2_VGRID
@@ -24,12 +24,12 @@ def test_zco_orca2():
     # Bathymetry dataset
     ds_bathy = Bathymetry(1.0e3, 1.2e3, 1, 1).flat(5.0e3)
 
-    # zco grid generator
-    zco = Zco(ds_bathy, jpk=31)
+    # Set number of vertical levels
+    ds_bathy.domcfg.jpk = 31
 
     # zco mesh with analytical e3 using ORCA2 input parameters
     # See pag 62 of v3.6 manual for the input parameters
-    dsz_an = zco(
+    dsz_an = ds_bathy.domcfg.zco(
         ppdzmin=10.0,
         pphmax=5000.0,
         ppkth=21.43336197938,
@@ -68,12 +68,12 @@ def test_zco_uniform():
     # Bathymetry dataset
     ds_bathy = Bathymetry(10.0e3, 1.2e3, 1, 1).flat(5.0e3)
 
-    # zco grid generator
-    zco = Zco(ds_bathy, jpk=51)
+    # Set number of vertical levels
+    ds_bathy.domcfg.jpk = 31
 
     # Compare zco mesh with analytical VS finite difference e3
-    expected = zco(**kwargs, ln_e3_dep=True)
-    actual = zco(**kwargs, ln_e3_dep=False)
+    expected = ds_bathy.domcfg.zco(**kwargs, ln_e3_dep=True)
+    actual = ds_bathy.domcfg.zco(**kwargs, ln_e3_dep=False)
     eps = 1.0e-14  # truncation errors
     xr.testing.assert_allclose(expected, actual, rtol=eps, atol=0)
 
@@ -83,8 +83,8 @@ def test_zco_x_y_invariant():
 
     # Generate 2x2 flat bathymetry dataset
     ds_bathy = Bathymetry(10.0e3, 1.2e3, 2, 2).flat(5.0e3)
-    zco = Zco(ds_bathy, jpk=10)
-    ds = zco(ppdzmin=10, pphmax=5.0e3)
+    ds_bathy.domcfg.jpk = 10
+    ds = ds_bathy.domcfg.zco(ppdzmin=10, pphmax=5.0e3)
 
     # Check z3 and e3
     for varname in ["z3T", "z3W", "e3T", "e3W"]:
@@ -101,50 +101,50 @@ def test_zco_errors():
 
     # Generate test data
     ds_bathy = Bathymetry(1.0e3, 1.2e3, 1, 1).flat(5.0e3)
-    zco = Zco(ds_bathy, jpk=10)
+    ds_bathy.domcfg.jpk = 10
 
     # Without ldbletanh flag, only allow all pps set or none of them
     with pytest.raises(
         ValueError, match="ppa2, ppkth2 and ppacr2 MUST be all None or all float"
     ):
-        zco(**kwargs, ppa2=1, ppkth2=1, ppacr2=None)
+        ds_bathy.domcfg.zco(**kwargs, ppa2=1, ppkth2=1, ppacr2=None)
 
     # When ldbletanh flag is True, all coefficients must be specified
     with pytest.raises(ValueError, match="ppa2, ppkth2 and ppacr2 MUST be all float"):
-        zco(**kwargs, ldbletanh=True, ppa2=1, ppkth2=1, ppacr2=None)
+        ds_bathy.domcfg.zco(**kwargs, ldbletanh=True, ppa2=1, ppkth2=1, ppacr2=None)
 
 
 def test_zco_warnings():
     """Make sure we warn when arguments are ignored"""
 
     # Initialize test class
     ds_bathy = Bathymetry(1.0e3, 1.2e3, 1, 1).flat(5.0e3)
-    zco = Zco(ds_bathy, jpk=10)
+    ds_bathy.domcfg.jpk = 10
 
     # Uniform: Ignore stretching
     kwargs = dict(ppdzmin=10, pphmax=5.0e3, ppkth=0, ppacr=0)
-    expected = zco(**kwargs, ppsur=None, ppa0=999_999, ppa1=None)
+    expected = ds_bathy.domcfg.zco(**kwargs, ppsur=None, ppa0=999_999, ppa1=None)
     with pytest.warns(
         UserWarning, match="ppsur, ppa0 and ppa1 are ignored when ppacr == ppkth == 0"
     ):
-        actual = zco(**kwargs, ppsur=2, ppa0=2, ppa1=2)
+        actual = ds_bathy.domcfg.zco(**kwargs, ppsur=2, ppa0=2, ppa1=2)
     xr.testing.assert_identical(expected, actual)
 
     # ldbletanh OFF: Ignore double tanh
     kwargs = dict(ppdzmin=10, pphmax=5.0e3, ldbletanh=False)
-    expected = zco(**kwargs, ppa2=None, ppkth2=None, ppacr2=None)
+    expected = ds_bathy.domcfg.zco(**kwargs, ppa2=None, ppkth2=None, ppacr2=None)
     with pytest.warns(
         UserWarning, match="ppa2, ppkth2 and ppacr2 are ignored when ldbletanh is False"
     ):
-        actual = zco(**kwargs, ppa2=2, ppkth2=2, ppacr2=2)
+        actual = ds_bathy.domcfg.zco(**kwargs, ppa2=2, ppkth2=2, ppacr2=2)
     xr.testing.assert_identical(expected, actual)
 
     # Uniform case: Ignore double tanh
     kwargs = dict(ppdzmin=10, pphmax=5.0e3, ppkth=0, ppacr=0)
-    expected = zco(**kwargs, ppa2=None, ppkth2=None, ppacr2=None)
+    expected = ds_bathy.domcfg.zco(**kwargs, ppa2=None, ppkth2=None, ppacr2=None)
     with pytest.warns(
         UserWarning,
         match="ppa2, ppkth2 and ppacr2 are ignored when ppacr == ppkth == 0",
     ):
-        actual = zco(**kwargs, ppa2=2, ppkth2=2, ppacr2=2)
+        actual = ds_bathy.domcfg.zco(**kwargs, ppa2=2, ppkth2=2, ppacr2=2)
     xr.testing.assert_identical(expected, actual)