Fix tree printer to include all object types

doc/source/user_guide/howto/print_model.rst

@@ -21,28 +21,27 @@ You can print the tree structure using the :func:`.print_model` function:
 .. doctest::
 
     >>> pyacp.print_model(model)
-    Model
-        Material Data
-            Materials
-                Structural Steel
-            Fabrics
-                Fabric.1
+    'ACP Model'
+        Materials
+            'Structural Steel'
+        Fabrics
+            'Fabric.1'
         Element Sets
-            All_Elements
+            'All_Elements'
         Edge Sets
-            ns_edge
-        Geometry
+            'ns_edge'
         Rosettes
-            Global Coordinate System
-        Lookup Tables
-        Selection Rules
+            'Global Coordinate System'
         Oriented Selection Sets
-            OrientedSelectionSet.1
+            'OrientedSelectionSet.1'
         Modeling Groups
-            ModelingGroup.1
-                ModelingPly.1
-                    ProductionPly
-                        P1L1__ModelingPly.1
+            'ModelingGroup.1'
+                Modeling Plies
+                    'ModelingPly.1'
+                        Production Plies
+                            'P1__ModelingPly.1'
+                                Analysis Plies
+                                    'P1L1__ModelingPly.1'
     <BLANKLINE>
 
 
@@ -52,16 +51,66 @@ Alternatively, you can use :func:`.get_model_tree` to get a tree representation.
 
     >>> tree_root = pyacp.get_model_tree(model)
     >>> tree_root.label
-    'Model'
+    "'ACP Model'"
     >>> for child in tree_root.children:
     ...     print(child.label)
     ...
-    Material Data
+    Materials
+    Fabrics
     Element Sets
     Edge Sets
-    Geometry
     Rosettes
-    Lookup Tables
-    Selection Rules
     Oriented Selection Sets
     Modeling Groups
+
+
+The ``hide_empty`` label can be set to ``False`` to also show empty groups:
+
+.. doctest::
+
+    >>> pyacp.print_model(model, hide_empty=False)
+    'ACP Model'
+        Materials
+            'Structural Steel'
+        Fabrics
+            'Fabric.1'
+        Stackups
+        Sublaminates
+        Element Sets
+            'All_Elements'
+        Edge Sets
+            'ns_edge'
+        Cad Geometries
+        Virtual Geometries
+        Rosettes
+            'Global Coordinate System'
+        Lookup Tables 1d
+        Lookup Tables 3d
+        Parallel Selection Rules
+        Cylindrical Selection Rules
+        Spherical Selection Rules
+        Tube Selection Rules
+        Cutoff Selection Rules
+        Geometrical Selection Rules
+        Variable Offset Selection Rules
+        Boolean Selection Rules
+        Oriented Selection Sets
+            'OrientedSelectionSet.1'
+        Modeling Groups
+            'ModelingGroup.1'
+                Modeling Plies
+                    'ModelingPly.1'
+                        Production Plies
+                            'P1__ModelingPly.1'
+                                Analysis Plies
+                                    'P1L1__ModelingPly.1'
+                Interface Layers
+                Butt Joint Sequences
+        Imported Modeling Groups
+        Sampling Points
+        Section Cuts
+        Solid Models
+        Imported Solid Models
+        Sensors
+        Field Definitions
+    <BLANKLINE>

src/ansys/acp/core/_model_printer.py

@@ -22,6 +22,8 @@
 
 import os
 
+from ._tree_objects._grpc_helpers.mapping import Mapping
+from ._tree_objects.base import TreeObjectBase
 from ._tree_objects.model import Model
 from ._utils.string_manipulation import replace_underscores_and_capitalize
 
@@ -52,95 +54,56 @@
         return ret
 
 
-def _add_tree_part(
-    tree: Node,
-    container_name: str,
-    model: Model,
-) -> None:
-    items = list(getattr(model, container_name).items())
-    if len(items) == 0:
-        return
-    container = Node(replace_underscores_and_capitalize(container_name))
-    tree.children.append(container)
-    for entity_name, entity in items:
-        group_node = Node(entity_name)
-        container.children.append(group_node)
-
-
-def print_model(model: Model) -> None:
+def print_model(model: Model, *, hide_empty: bool = True) -> None:
     """Print a tree representation of the model.
 
     Parameters
     ----------
     model:
         pyACP model
+    hide_empty :
+        Whether to hide empty collections.
 
     """
-    return print(get_model_tree(model))
+    return print(get_model_tree(model, hide_empty=hide_empty))
 
 
-def get_model_tree(model: Model) -> Node:
+def get_model_tree(model: Model, *, hide_empty: bool = True) -> Node:
     """Get a tree representation of the model.
 
     Returns the root node.
 
     Parameters
     ----------
-    model:
-        pyACP model.
+    model :
+        ACP model.
+    hide_empty :
+        Whether to hide empty collections.
     """
-    model_node = Node("Model")
-
-    material_data = Node("Material Data")
-    model_node.children.append(material_data)
-    _add_tree_part(material_data, "materials", model)
-    _add_tree_part(material_data, "fabrics", model)
-    _add_tree_part(material_data, "stackups", model)
-    _add_tree_part(material_data, "sublaminates", model)
-
-    _add_tree_part(model_node, "element_sets", model)
-    _add_tree_part(model_node, "edge_sets", model)
-
-    geometry = Node("Geometry")
-    model_node.children.append(geometry)
-    _add_tree_part(geometry, "cad_geometries", model)
-    _add_tree_part(geometry, "virtual_geometries", model)
-
-    _add_tree_part(model_node, "rosettes", model)
-
-    lookup_table = Node("Lookup Tables")
-    model_node.children.append(lookup_table)
-    _add_tree_part(lookup_table, "lookup_tables_1d", model)
-    _add_tree_part(lookup_table, "lookup_tables_3d", model)
-
-    selection_rules = Node("Selection Rules")
-    model_node.children.append(selection_rules)
-    _add_tree_part(selection_rules, "parallel_selection_rules", model)
-    _add_tree_part(selection_rules, "cylindrical_selection_rules", model)
-    _add_tree_part(selection_rules, "spherical_selection_rules", model)
-    _add_tree_part(selection_rules, "tube_selection_rules", model)
-    _add_tree_part(selection_rules, "cutoff_selection_rules", model)
-    _add_tree_part(selection_rules, "geometrical_selection_rules", model)
-    _add_tree_part(selection_rules, "variable_offset_selection_rules", model)
-    _add_tree_part(selection_rules, "boolean_selection_rules", model)
-
-    _add_tree_part(model_node, "oriented_selection_sets", model)
-
-    modeling_groups = Node("Modeling Groups")
-    model_node.children.append(modeling_groups)
-    for modeling_group_name, modeling_group in model.modeling_groups.items():
-        group_node = Node(modeling_group_name)
-        modeling_groups.children.append(group_node)
-        for modeling_ply_name, modeling_ply in modeling_group.modeling_plies.items():
-            modeling_ply_node = Node(modeling_ply_name)
-            group_node.children.append(modeling_ply_node)
-            for production_ply_name, production_ply in modeling_ply.production_plies.items():
-                production_ply_node = Node(production_ply_name)
-                modeling_ply_node.children.append(production_ply_node)
-                for analysis_ply_name, analysis_ply in production_ply.analysis_plies.items():
-                    analysis_ply_node = Node(analysis_ply_name)
-                    production_ply_node.children.append(analysis_ply_node)
-
-    _add_tree_part(model_node, "sensors", model)
-
-    return model_node
+    return _get_model_tree_impl(obj=model, hide_empty=hide_empty)
+
+
+def _get_model_tree_impl(obj: TreeObjectBase, *, hide_empty: bool) -> Node:
+    obj_node = Node(repr(_name_or_id(obj)))
+    for attr_name in obj._GRPC_PROPERTIES:
+        try:
+            attr = getattr(obj, attr_name)
+        except (AttributeError, RuntimeError):
+            continue
+        if isinstance(attr, Mapping):
+            collection_node = Node(replace_underscores_and_capitalize(attr_name))
+            obj_node.children.append(collection_node)
+            for child_obj in attr.values():
+                collection_node.children.append(
+                    _get_model_tree_impl(child_obj, hide_empty=hide_empty)
+                )
+            if hide_empty and not collection_node.children:
+                obj_node.children.pop()
+    return obj_node
+
+
+def _name_or_id(obj: TreeObjectBase) -> str:
+    try:
+        return obj.name
+    except AttributeError:
+        return obj.id  # type: ignore