refactor(createpackages): use jinja for mf6 module code generation

autotest/test_codegen.py

@@ -8,9 +8,7 @@
 MF6_PATH = PROJ_ROOT / "flopy" / "mf6"
 DFN_PATH = MF6_PATH / "data" / "dfn"
 DFN_NAMES = [
-    dfn.stem
-    for dfn in DFN_PATH.glob("*.dfn")
-    if dfn.stem not in ["common", "flopy"]
+    dfn.stem for dfn in DFN_PATH.glob("*.dfn") if dfn.stem not in ["common", "flopy"]
 ]
 
 

flopy/mf6/utils/codegen/__init__.py

@@ -12,7 +12,13 @@ def _get_template_env():
     # expect optional deps at module init time
     jinja = import_optional_dependency("jinja2")
     loader = jinja.PackageLoader("flopy", "mf6/utils/codegen/templates/")
-    env = jinja.Environment(loader=loader)
+    env = jinja.Environment(
+        loader=loader,
+        trim_blocks=True,
+        lstrip_blocks=True,
+        line_statement_prefix="_",
+        keep_trailing_newline=True,
+    )
 
     from flopy.mf6.utils.codegen.filters import Filters
 
@@ -47,9 +53,7 @@ def make_init(dfns: dict, outdir: PathLike, verbose: bool = False):
     from flopy.mf6.utils.codegen.context import Context
 
     contexts = list(
-        chain(
-            *[[ctx for ctx in Context.from_dfn(dfn)] for dfn in dfns.values()]
-        )
+        chain.from_iterable(Context.from_dfn(dfn) for dfn in dfns.values())
     )
     target_name = "__init__.py"
     target_path = outdir / target_name

flopy/mf6/utils/codegen/context.py

@@ -69,8 +69,8 @@ def from_dfn(dfn: Dfn) -> List["Context.Name"]:
     name: Name
     vars: Vars
 
-    @classmethod
-    def from_dfn(cls, dfn: Dfn) -> Iterator["Context"]:
+    @staticmethod
+    def from_dfn(dfn: Dfn) -> Iterator["Context"]:
         """
         Extract context class descriptor(s) from an input definition.
         These are structured representations of input context classes.

flopy/mf6/utils/codegen/dfn.py

@@ -396,7 +396,7 @@ def _fields() -> Vars:
                     block=block,
                     children=None,
                     description=(
-                        f"* Contains data for the {ref['abbr']} package. Data can be "
+                        f"Contains data for the {ref['abbr']} package. Data can be "
                         f"stored in a dictionary containing data for the {ref['abbr']} "
                         "package with variable names as keys and package data as "
                         f"values. Data just for the {ref['val']} variable is also "

flopy/mf6/utils/codegen/filters.py

@@ -82,10 +82,6 @@ def parent(ctx, ctx_name) -> str:
                 or ctx_name.l in ["sim", "exg", "sln"]
             ):
                 return "simulation"
-            elif subpkg:
-                if ctx_name.l == "utl" and ctx_name.r == "hpc":
-                    return "simulation"
-                return "package"
             return "model"
 
         @pass_context
@@ -101,13 +97,7 @@ def skip_init(ctx, ctx_name) -> List[str]:
                     "solutiongroup",
                 ]
             elif base == "MFModel":
-                skip = ["packages", "export_netcdf", "nc_filerecord"]
-                refs = ctx.get("foreign_keys", dict())
-                if any(refs) and ctx["name"] != (None, "nam"):
-                    for k in refs.keys():
-                        if ctx["vars"].get(k, None):
-                            skip.append(k)
-                return skip
+                return ["packages", "export_netcdf", "nc_filerecord"]
             else:
                 if ctx_name.r == "nam":
                     return ["export_netcdf", "nc_filerecord"]
@@ -156,17 +146,17 @@ def type(var: dict) -> str:
             """
             Get a readable representation of the variable's type.
             TODO: eventually replace this with a proper `type` in
-            the variable spec when we add type hints. For now try
-            to match the existing format, with a few tweaks; e.g.
-            distinguishing lists from records by square and round
-            brackets, respectively, and separating each choice in
-            a keystring by '|'.
+            the variable spec when we add type hints
             """
             _type = var["type"]
             shape = var.get("shape", None)
             children = var.get("children", None)
             if children:
                 if _type == "list":
+                    if len(children) == 1:
+                        first = list(children.values())[0]
+                        if first["type"] in ["record", "union"]:
+                            return f"[{Filters.Var.type(first)}]"
                     children = ", ".join(
                         [v["name"] for v in children.values()]
                     )
@@ -295,9 +285,7 @@ def _var(var: dict) -> List[str]:
 
             attrs = list(filter(None, [_attr(v) for v in variables.values()]))
 
-            if base == "MFModel":
-                attrs.append(f"model_type = {name.l}")
-            elif base == "MFPackage":
+            if base == "MFPackage":
                 attrs.extend(
                     [
                         f"package_abbr = '{name.r}'"

flopy/mf6/utils/codegen/templates/exchange.py.jinja

@@ -13,12 +13,12 @@ class Modflow{{ title }}(MFPackage):
 
     Parameters
     ----------
-    {{ macros.vars_docs(vars, indent=4) }}
+{{ macros.vars_docs(vars, indent=4) }}
     """
 
     {% for attr in vars|attrs %}
     {{ attr }}
-    {%- endfor %}
+    {% endfor %}
 
     def __init__(
         self,
@@ -27,7 +27,7 @@ class Modflow{{ title }}(MFPackage):
         exgtype="{{ name.r[:3].upper() }}6-{{ name.r[3:].upper() }}6",
         exgmnamea=None,
         exgmnameb=None,
-        {{ macros.init_vars(vars, indent=8, skip=name|skip_init) }}
+{{ macros.init_vars(vars, indent=8, skip=name|skip_init) }}
         filename=None,
         pname=None,
         **kwargs,
@@ -64,7 +64,7 @@ class Modflow{{ title }}(MFPackage):
             GWE Model with the name exgmnameb must correspond to the GWF Model
             with the name gwfmodelname2. 
 
-        {{ macros.vars_docs(vars, indent=8) }}
+{{ macros.vars_docs(vars, indent=8) }}
         """
 
         super().__init__(
@@ -83,6 +83,6 @@ class Modflow{{ title }}(MFPackage):
 
         {% for statement in vars|init %}
         {{ statement }}
-        {%- endfor %}
+        {% endfor %}
 
         self._init_complete = True

flopy/mf6/utils/codegen/templates/macros.jinja

@@ -1,21 +1,25 @@
 {% macro init_vars(vars, alias=false, indent=0, skip=none) %}
-{%- for name, var in vars.items() if name not in skip %}
-{% set v = var|untag -%}
-{% set n = (name if alias else v.name)|safe_name -%}
-{{ ""|indent(indent, first=true) }}{{ n }}{%- if v.default is defined %}={{ v.default|value }}{%- endif -%},
-{%- endfor %}
+{% for name, var in vars.items() if name not in skip %}
+{% set v = var|untag %}
+{% set n = (name if alias else v.name)|safe_name %}
+{{ ""|indent(indent, first=true) }}{{ n }}{% if v.default is defined %}={{ v.default|value }}{% endif %},
+{% endfor %}
 {% endmacro %}
 
-{% macro vars_docs(vars, indent=0) %}
-{%- for var in vars.values() recursive %}
-{% set v = var|untag -%}
-{% set n = v.name|safe_name|escape_trailing_underscore -%}
+{% macro vars_docs(vars, indent=0, recurse=true) %}
+{% for var in vars.values() recursive %}
+{% set v = var|untag %}
+{% set n = v.name|safe_name|escape_trailing_underscore %}
 {{ ""|indent(indent, first=true) }}{% if loop.depth > 1 %}* {% endif %}{{ n }} : {{ v|type }}
-{%- if v.description is defined and v.description is not none %}
+{% if v.description is defined and v.description is not none %}
 {{ v.description|wordwrap|indent(indent + (loop.depth * 4), first=true) }}
-{%- endif %}
-{%- if v.children is defined and v.children is not none -%}
+{% endif %}
+{% if recurse and v.children is defined and v.children is not none %}
+{% if v.type == "list" and v.children|length == 1 and (v.children.values()|first).type == "record" %}
+{{ loop((v.children.values()|first).children.values())|indent(indent, first=true) }} 
+{% else %}
 {{ loop(v.children.values())|indent(indent, first=true) }}
-{%- endif %}
-{% endfor -%}
-{% endmacro %}
+{% endif %}
+{% endif %}
+{% endfor %}
+{% endmacro %}

flopy/mf6/utils/codegen/templates/model.py.jinja

@@ -14,7 +14,7 @@ class Modflow{{ title }}(MFModel):
 
     Parameters
     ----------
-    {{ macros.vars_docs(vars, indent=4) }}
+{{ macros.vars_docs(vars, indent=4) }}
 
     Methods
     -------
@@ -34,7 +34,7 @@ class Modflow{{ title }}(MFModel):
         version="mf6",
         exe_name="mf6",
         model_rel_path=".",
-        {{ macros.init_vars(vars, indent=8, skip=name|skip_init) }}
+{{ macros.init_vars(vars, indent=8, skip=name|skip_init) }}
         **kwargs,
     ):
         """
@@ -61,7 +61,7 @@ class Modflow{{ title }}(MFModel):
             Simulation that this model is a part of.  Model is automatically
             added to simulation when it is initialized.
 
-        {{ macros.vars_docs(vars, indent=8) }}
+{{ macros.vars_docs(vars, indent=8) }}
         """
 
         super().__init__(
@@ -77,7 +77,7 @@ class Modflow{{ title }}(MFModel):
 
         {% for statement in vars|init %}
         {{ statement }}
-        {%- endfor %}
+        {% endfor %}
 
     @classmethod
     def load(

flopy/mf6/utils/codegen/templates/package.py.jinja

@@ -14,18 +14,18 @@ class Modflow{{ title }}(MFPackage):
 
     Parameters
     ----------
-    {{ macros.vars_docs(vars, indent=4) }}
+{{ macros.vars_docs(vars, indent=4) }}
     """
 
     {% for attr in vars|attrs %}
     {{ attr }}
-    {%- endfor %}
+    {% endfor %}
 
     def __init__(
         self,
         {{ name|parent }},
         loading_package=False,
-        {{ macros.init_vars(vars, indent=8, skip=name|skip_init) }}
+{{ macros.init_vars(vars, indent=8, skip=name|skip_init) }}
         filename=None,
         pname=None,
         **kwargs,
@@ -43,7 +43,7 @@ class Modflow{{ title }}(MFPackage):
             Do not set this parameter. It is intended for debugging and internal
             processing purposes only.
 
-        {{ macros.vars_docs(vars, indent=8) }}
+{{ macros.vars_docs(vars, indent=8) }}
 
         filename : str
             File name for this package.
@@ -68,7 +68,7 @@ class Modflow{{ title }}(MFPackage):
 
         {% for statement in vars|init %}
         {{ statement }}
-        {%- endfor %}
+        {% endfor %}
 
         self._init_complete = True
 
@@ -94,15 +94,15 @@ class {{ title }}Packages(MFChildPackages):
 
     def initialize(
         self,
-        {{ macros.init_vars(vars, alias=true, indent=8, skip=name|skip_init) }}
+{{ macros.init_vars(vars, alias=true, indent=8, skip=name|skip_init) }}
         filename=None,
         pname=None,
     ):
         new_package = Modflow{{ title }}(
             self._cpparent,
-            {%- for n, var in vars.items() if n not in name|skip_init %}
+            {% for n, var in vars.items() if n not in name|skip_init %}
             {{ n|safe_name }}={{ n|safe_name }},
-            {%- endfor %}
+            {% endfor %}
             filename=filename,
             pname=pname,
             child_builder_call=True,
@@ -112,15 +112,15 @@ class {{ title }}Packages(MFChildPackages):
     {% if name.r != "obs" %}
     def append_package(
         self,
-        {{ macros.init_vars(vars, alias=true, indent=8, skip=name|skip_init) }}
+{{ macros.init_vars(vars, alias=true, indent=8, skip=name|skip_init) }}
         filename=None,
         pname=None,
     ):
         new_package = Modflow{{ title }}(
             self._cpparent,
-            {%- for n, var in vars.items() if n not in name|skip_init %}
+            {% for n, var in vars.items() if n not in name|skip_init %}
             {{ n|safe_name }}={{ n|safe_name }},
-            {%- endfor %}
+            {% endfor %}
             filename=filename,
             pname=pname,
             child_builder_call=True,

flopy/mf6/utils/codegen/templates/simulation.py.jinja

@@ -13,7 +13,7 @@ class MF{{ title }}(MFSimulationBase):
 
     Parameters
     ----------
-    {{ macros.vars_docs(vars, indent=4) }}
+{{ macros.vars_docs(vars, indent=4) }}
 
     Methods
     -------
@@ -35,7 +35,7 @@ class MF{{ title }}(MFSimulationBase):
         write_headers: bool = True,
         use_pandas: bool = True,
         lazy_io: bool = False,
-        {{ macros.init_vars(vars, indent=8, skip=name|skip_init) }}
+{{ macros.init_vars(vars, indent=8, skip=name|skip_init) }}
     ):
         """
         {{ name|description }}
@@ -66,7 +66,7 @@ class MF{{ title }}(MFSimulationBase):
         lazy_io
             Whether to use lazy IO
 
-        {{ macros.vars_docs(vars, indent=8) }}
+{{ macros.vars_docs(vars, indent=8) }}
         """
 
         super().__init__(
@@ -82,7 +82,7 @@ class MF{{ title }}(MFSimulationBase):
 
         {% for statement in vars|init %}
         {{ statement }}
-        {%- endfor %}
+        {% endfor %}
 
     @classmethod
     def load(