Implement accessor (#36)

* Implement accessor

* define a decorator to handle errors

* Better decorator and error handling

* Update pydomcfg/accessor.py

* Update pydomcfg/accessor.py

* Update pydomcfg/accessor.py

* cast wrapper

* fix docs

* add tests

* what's new

* Add attributes to api

* bloody darglint

Author: malmans2

ci/docs.yml

@@ -4,5 +4,6 @@ channels:
 dependencies:
   - numpydoc
   - sphinx
+  - sphinx-autosummary-accessors
   - sphinx_rtd_theme
   - xarray

docs/conf.py

@@ -9,6 +9,7 @@
 import os
 import sys
 
+import sphinx_autosummary_accessors
 import sphinx_rtd_theme  # noqa: F401
 
 # isort: off
@@ -40,8 +41,10 @@
     "sphinx.ext.autosummary",
     "sphinx.ext.viewcode",
     "sphinx.ext.extlinks",
+    "sphinx.ext.intersphinx",
     "sphinx_rtd_theme",
     "numpydoc",
+    "sphinx_autosummary_accessors",
 ]
 
 # GitHub links
@@ -50,7 +53,7 @@
 }
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ["_templates"]
+templates_path = ["_templates", sphinx_autosummary_accessors.templates_path]
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.

docs/developers/whats-new.rst

@@ -6,7 +6,8 @@ What's New
 Unreleased
 ==========
 
-- Added :py:func:`~pydomcfg.bathymetry.sea_mount`
+- Introduced the ``domcfg`` accessor (:pr:`36`)
+- Added :py:meth:`~pydomcfg.tests.bathymetry.Bathymetry.sea_mount`
   useful to generate classic sea mount test case. (:pr:`17`)
 - Added :py:class:`~pydomcfg.domzgr.zco.Zco`
   to generate geopotential z-coordinates. (:pr:`15`)

docs/users/api.rst

@@ -1,12 +1,31 @@
-.. currentmodule:: pydomcfg
-
 API reference
 #############
 
 This page provides an auto-generated summary of pyDOMCFG's API.
 
+Accessor
+========
+.. currentmodule:: xarray
+
+Attributes
+----------
+.. autosummary::
+   :toctree: generated/
+   :template: autosummary/accessor_attribute.rst
+
+    Dataset.domcfg.jpk
+
+Methods
+-------
+.. autosummary::
+   :toctree: generated/
+   :template: autosummary/accessor_method.rst
+
+    Dataset.domcfg.zco
+
 Domzgr
 ======
+.. currentmodule:: pydomcfg
 
 Zco
 ----
@@ -18,6 +37,8 @@ Zco
 
 Utils
 =====
+.. currentmodule:: pydomcfg
+
 .. autosummary::
    :toctree: generated/
 

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,54 @@
+from typing import Any, Callable, TypeVar, cast
+
+import xarray as xr
+from xarray import Dataset
+
+from .domzgr.zco import Zco
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def _jpk_check(func: F) -> F:
+    """
+    Decorator to raise an error if jpk was not set
+    """
+
+    def wrapper(self, *args, **kwargs):
+        if not self.jpk:
+            raise ValueError(
+                f"Set `jpk` before calling `obj.domcfg.{func.__name__}`."
+                " For example: obj.domcfg.jpk = 31"
+            )
+        return func(self, *args, **kwargs)
+
+    return cast(F, wrapper)
+
+
+@xr.register_dataset_accessor("domcfg")
+class Accessor:
+    def __init__(self, xarray_obj: Dataset):
+        self._obj = xarray_obj
+        self._jpk = 0
+
+    @property
+    def jpk(self) -> int:
+        """
+        Number of computational levels
+
+        Returns
+        -------
+        int
+        """
+        return self._jpk
+
+    @jpk.setter
+    def jpk(self, value: int):
+        if value <= 0:
+            raise ValueError("`jpk` MUST be > 0")
+        self._jpk = value
+
+    @_jpk_check
+    def zco(self, *args: Any, **kwargs: Any) -> Dataset:
+        return Zco(self._obj, self._jpk)(*args, **kwargs)
+
+    zco.__doc__ = Zco.__call__.__doc__

pydomcfg/domzgr/zco.py

@@ -88,6 +88,7 @@ def __call__(
 
         See Also
         --------
+        pydomcfg.domzgr.zco.Zco.__call__: Underlying method.
         NEMOv4.0: ``domzgr/zgr_z`` subroutine
 
         References

pydomcfg/tests/test_accessor.py

@@ -0,0 +1,24 @@
+import pytest
+
+import pydomcfg  # noqa: F401
+
+from .bathymetry import Bathymetry
+
+
+def test_jpk():
+    """Make sure jpk is set correctly"""
+
+    # Initialize test class
+    ds_bathy = Bathymetry(1.0e3, 1.2e3, 1, 1).flat(5.0e3)
+
+    # Can't call methods without setting jpk first
+    with pytest.raises(ValueError, match="Set `jpk` before calling `obj.domcfg.zco`"):
+        ds_bathy.domcfg.zco()
+
+    # jpk must be > 0
+    with pytest.raises(ValueError, match="`jpk` MUST be > 0"):
+        ds_bathy.domcfg.jpk = -1
+
+    # Has been set correctly.
+    ds_bathy.domcfg.jpk = 1
+    assert ds_bathy.domcfg.jpk == 1

pydomcfg/tests/test_zco.py

@@ -6,7 +6,7 @@
 import pytest
 import xarray as xr
 
-from pydomcfg.domzgr.zco import Zco
+import pydomcfg  # noqa: F401
 
 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)