1- from typing import Any , Callable , TypeVar , cast
1+ import inspect
2+ import warnings
3+ from collections import ChainMap
4+ from functools import wraps
5+ from pathlib import Path
6+ from typing import IO , Any , Callable , TypeVar , Union , cast
27
38import xarray as xr
49from xarray import Dataset
510
611from .domzgr .zco import Zco
712
13+ try :
14+ import f90nml
15+
16+ HAS_F90NML = True
17+ except ImportError :
18+ HAS_F90NML = False
19+
820F = TypeVar ("F" , bound = Callable [..., Any ])
21+ ZGR_MAPPER = {
22+ "ln_zco" : Zco ,
23+ # "ln_zps": TODO,
24+ # "ln_sco": TODO
25+ }
926
1027
1128def _jpk_check (func : F ) -> F :
1229 """
1330 Decorator to raise an error if jpk was not set
1431 """
1532
33+ @wraps (func )
1634 def wrapper (self , * args , ** kwargs ):
1735 if not self .jpk :
1836 raise ValueError (
@@ -29,7 +47,9 @@ class Accessor:
2947 def __init__ (self , xarray_obj : Dataset ):
3048 self ._obj = xarray_obj
3149 self ._jpk = 0
50+ self ._nml_ref_path = ""
3251
52+ # Set attributes
3353 @property
3454 def jpk (self ) -> int :
3555 """
@@ -43,12 +63,112 @@ def jpk(self) -> int:
4363
4464 @jpk .setter
4565 def jpk (self , value : int ):
46- if value <= 0 :
47- raise ValueError ("`jpk` MUST be > 0 " )
66+ if value < 0 :
67+ raise ValueError ("`jpk` MUST be >= 0 (use 0 to unset jpk) " )
4868 self ._jpk = value
4969
70+ @property
71+ def nml_ref_path (self ) -> str :
72+ """
73+ Path to reference namelist.
74+
75+ Returns
76+ -------
77+ str
78+ """
79+ return self ._nml_ref_path
80+
81+ @nml_ref_path .setter
82+ def nml_ref_path (self , value : str ):
83+ self ._nml_ref_path = value
84+
85+ # domzgr methods
86+ # TODO:
87+ # I think the process of creating the public API and doc
88+ # can be further automatized, but let's not put too much effort into it
89+ # until we settle on the back-end structure:
90+ # See: https://github.com/pyNEMO/pyDOMCFG/issues/45
5091 @_jpk_check
5192 def zco (self , * args : Any , ** kwargs : Any ) -> Dataset :
52- return Zco (self ._obj , self ._jpk )(* args , ** kwargs )
93+ name = inspect .stack ()[0 ][3 ]
94+ return ZGR_MAPPER ["ln_" + name ](self ._obj , self ._jpk )(* args , ** kwargs )
5395
5496 zco .__doc__ = Zco .__call__ .__doc__
97+
98+ # Emulate NEMO DOMAINcfg tools
99+ def from_namelist (self , nml_cfg_path_or_io : Union [str , Path , IO [str ]]) -> Dataset :
100+ """
101+ Auto-populate pydomcfg parameters using NEMO DOMAINcfg namelists.
102+
103+ Parameters
104+ ----------
105+ nml_cfg_path_or_io: str, Path, IO
106+ Path pointing to a namelist_cfg,
107+ or namelist_cfg previously opened with open()
108+
109+ Returns
110+ -------
111+ Dataset
112+ """
113+
114+ nml_chained = self ._namelist_parser (nml_cfg_path_or_io )
115+ zgr_initialized , kwargs = self ._get_zgr_initialized_and_kwargs (nml_chained )
116+ return zgr_initialized (** kwargs )
117+
118+ def _namelist_parser (
119+ self , nml_cfg_path_or_io : Union [str , Path , IO [str ]]
120+ ) -> ChainMap :
121+ """Parse namelists using f90nml, chaining all namblocks"""
122+
123+ if not HAS_F90NML :
124+ raise ImportError (
125+ "`f90nml` MUST be installed to use `obj.domcfg.from_namelist()`"
126+ )
127+
128+ if not self .nml_ref_path :
129+ raise ValueError (
130+ "Set `nml_ref_path` before calling `obj.domcfg.from_namelist()`"
131+ " For example: obj.domcfg.nml_ref_path = 'path/to/nml_ref'"
132+ )
133+
134+ if self .jpk :
135+ warnings .warn (
136+ "`obj.domcfg.jpk` is ignored. `jpk` is inferred from the namelists."
137+ )
138+
139+ # Read namelists: cfg overrides ref
140+ nml_cfg = f90nml .read (nml_cfg_path_or_io )
141+ nml = f90nml .patch (self .nml_ref_path , nml_cfg )
142+
143+ return ChainMap (* nml .todict ().values ())
144+
145+ def _get_zgr_initialized_and_kwargs (self , nml_chained : ChainMap ):
146+
147+ # TODO: Add return type hint when abstraction in base class is implemented
148+
149+ # Pick the appropriate class
150+ zgr_classes = [
151+ value for key , value in ZGR_MAPPER .items () if nml_chained .get (key )
152+ ]
153+ if len (zgr_classes ) != 1 :
154+ raise ValueError (
155+ "One and only one of the following variables MUST be `.true.`:"
156+ f" { tuple (ZGR_MAPPER )}
157+ )
158+ zgr_class = zgr_classes [0 ]
159+
160+ # Compatibility with NEMO DOMAINcfg
161+ if nml_chained .get ("ldbletanh" ) is False :
162+ for pp in ["ppa2" , "ppkth2" , "ppacr2" ]:
163+ nml_chained [pp ] = None
164+
165+ # Get kwargs, converting 999_999 to None
166+ parameters = list (inspect .signature (zgr_class .__call__ ).parameters )
167+ parameters .remove ("self" )
168+ kwargs = {
169+ key : None if nml_chained [key ] == 999_999 else nml_chained [key ]
170+ for key in parameters
171+ if key in nml_chained
172+ }
173+
174+ return zgr_class (self ._obj , nml_chained ["jpkdta" ]), kwargs
0 commit comments