From 3596af5d8278d1bf7818818a3d63eae45a21458a Mon Sep 17 00:00:00 2001
From: wpbonelli <wbonelli@ucar.edu>
Date: Tue, 15 Oct 2024 18:28:10 -0400
Subject: [PATCH] docstring fixes

---
 flopy/mf6/utils/codegen/dfn.py                | 32 ++++++++++---------
 .../codegen/templates/docstring_params.jinja  |  2 +-
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/flopy/mf6/utils/codegen/dfn.py b/flopy/mf6/utils/codegen/dfn.py
index 667f31290..b30537f57 100644
--- a/flopy/mf6/utils/codegen/dfn.py
+++ b/flopy/mf6/utils/codegen/dfn.py
@@ -203,7 +203,7 @@ def load(
         referenced = dict()
         vars, meta = Dfn._load(f, **kwargs)
 
-        def _map(spec: Dict[str, Any], wrap: bool = False) -> Var:
+        def _map(spec: Dict[str, Any]) -> Var:
             """
             Convert a variable specification from its representation
             in an input definition file to a Pythonic form.
@@ -218,9 +218,6 @@ def _map(spec: Dict[str, Any], wrap: bool = False) -> Var:
             If a `default_value` is not provided, keywords are `False`
             by default, everything else is `None`.
 
-            If `wrap` is true, scalars will be wrapped as records.
-            This is useful to distinguish among choices in unions.
-
             Any filepath variable whose name functions as a foreign key
             for another context will be given a pointer to the context.
 
@@ -322,13 +319,17 @@ def _is_implicit_scalar_record():
                     children = {names[0]: _map(record)}
                     kind = Var.Kind.List
                 elif _is_implicit_scalar_record():
+                    fields = _fields(_name)
                     children = {
                         _name: Var(
                             name=_name,
                             kind=Var.Kind.Record,
                             block=block,
-                            children=_fields(_name),
+                            children=fields,
                             description=description,
+                            meta={
+                                "type": f"[{', '.join([f.meta["type"] for f in fields.values()])}]"
+                            },
                         )
                     }
                     kind = Var.Kind.List
@@ -349,31 +350,30 @@ def _is_implicit_scalar_record():
                             block=block,
                             children=first.children if single else fields,
                             description=description,
+                            meta={
+                                "type": f"[{', '.join([v.meta["type"] for v in fields.values()])}]"
+                            },
                         )
                     }
                     kind = Var.Kind.List
+                type_ = f"[{', '.join([v.name for v in children.values()])}]"
 
-            # union (product), children are choices.
-            # scalar choices are wrapped as records.
+            # union (product), children are choices
             elif _type.startswith("keystring"):
                 names = _type.split()[1:]
                 children = {
-                    v["name"]: _map(v, wrap=True)
+                    v["name"]: _map(v)
                     for v in vars.values(multi=True)
                     if v["name"] in names and v.get("in_record", False)
                 }
                 kind = Var.Kind.Union
+                type_ = f"[{', '.join([v.name for v in children.values()])}]"
 
             # record (sum), children are fields
             elif _type.startswith("record"):
                 children = _fields(_name)
                 kind = Var.Kind.Record
-
-            # are we wrapping a var into a record
-            # as a choice in a union?
-            elif wrap:
-                children = {_name: _map(spec)}
-                kind = Var.Kind.Record
+                type_ = f"[{', '.join([v.meta["type"] for v in children.values()])}]"
 
             # at this point, if it has a shape, it's an array
             elif shape is not None:
@@ -383,10 +383,12 @@ def _is_implicit_scalar_record():
                     kind = Var.Kind.List
                 else:
                     kind = Var.Kind.Array
+                type_ = f"[{_type}]"
 
             # finally scalars
             else:
                 kind = Var.Kind.Scalar
+                type_ = _type
 
             # create var
             return Var(
@@ -404,7 +406,7 @@ def _is_implicit_scalar_record():
                     else default
                 ),
                 children=children,
-                meta={"ref": ref},
+                meta={"ref": ref, "type": type_},
             )
 
         # pass the original DFN representation as
diff --git a/flopy/mf6/utils/codegen/templates/docstring_params.jinja b/flopy/mf6/utils/codegen/templates/docstring_params.jinja
index d8d5641d0..0b9d85d59 100644
--- a/flopy/mf6/utils/codegen/templates/docstring_params.jinja
+++ b/flopy/mf6/utils/codegen/templates/docstring_params.jinja
@@ -1,5 +1,5 @@
 {%- for v in vars.values() recursive %}
-    {% if loop.depth > 1 %}* {% endif %}{{ v.name }}{% if v._type is defined and v._type is not none %} : {{ v._type }}{% endif %}
+    {% if loop.depth > 1 %}* {% endif %}{{ v.name }}{% if v.meta is defined and v.meta.type is defined %} : {{ v.meta.type }}{% endif %}
 {%- if v.description is defined and v.description is not none %}
 {{ v.description|wordwrap|indent(4 + (loop.depth * 4), first=True) }}
 {%- endif %}