Merge pull request #101 from isteinbrecher/cmd_return_objects

Add functionality to get python objects created during string cmd

Author: isteinbrecher

src/cubitpy/cubitpy.py

@@ -29,7 +29,7 @@
 import netCDF4
 from fourcipp.fourc_input import FourCInput
 
-from cubitpy.conf import cupy
+from cubitpy.conf import GeometryType, cupy
 from cubitpy.cubit_group import CubitGroup
 from cubitpy.cubit_to_fourc_input import (
     add_exodus_geometry_section,
@@ -421,6 +421,87 @@ def reset(self):
         self.cubit.reset()
         self._default_cubit_variables()
 
+    def cmd_return(self, cmd: str, geometry_type: GeometryType, **kwargs):
+        """Run a cubit command and return the created geometry object.
+
+        Args:
+            cmd: The cubit command to run.
+            geometry_type: The geometry type that should be checked for a new geometry.
+            kwargs: Will be passed on to `cmd_return_dict`.
+
+        Returns:
+            If a single geometry object of the given type is created, this object
+            is returned. This function expects that a single geometry item of the given
+            type is created, otherwise an error will be raised. For use cases, where one
+            expects a variable amount of created items or wants to check multiple
+            different geometry types, please refer to `cmd_return_dict`.
+        """
+        geometry_dict = self.cmd_return_dict(cmd, [geometry_type], **kwargs)
+        if len(geometry_dict[geometry_type]) == 1:
+            return geometry_dict[geometry_type][0]
+        else:
+            raise ValueError(
+                f"Expected a single created item of type {geometry_type}, but got {geometry_dict[geometry_type]}"
+            )
+
+    def cmd_return_dict(
+        self,
+        cmd: str,
+        geometry_types: list[GeometryType],
+        *,
+        filter_sheet_bodies: bool = True,
+    ):
+        """Run a cubit command and return created geometry objects.
+
+        Args:
+            cmd: The cubit command to run.
+            geometry_types: The geometry types that should be checked for new geometries.
+            filter_sheet_bodies: If volumes that are sheet bodies should be ignored.
+                Defaults to true.
+
+        Returns:
+            A dictionary of the created geometry objects. The dictionary keys are the
+            geometry types, the values are lists containing the respective objects.
+        """
+
+        # Store the already existing ids for all requested geometry types.
+        geometry_ids_before = {
+            geometry: set(self.get_entities(geometry.get_cubit_string()))
+            for geometry in geometry_types
+        }
+
+        # For CoreForm, we need to check that the volumes are not sheet bodies
+        if cupy.is_coreform() and filter_sheet_bodies:
+            if cupy.geometry.volume in geometry_ids_before:
+                volume_ids_before = geometry_ids_before[cupy.geometry.volume]
+                volume_ids_before_no_sheet_bodies = {
+                    id for id in volume_ids_before if not self.is_sheet_body(id)
+                }
+                geometry_ids_before[cupy.geometry.volume] = (
+                    volume_ids_before_no_sheet_bodies
+                )
+
+        # Run the command.
+        self.cmd(cmd)
+
+        # Get the objects that were created by the command.
+        create_objects = {}
+        for geometry, ids_before in geometry_ids_before.items():
+            ids_after = set(self.get_entities(geometry.get_cubit_string()))
+            ids_new = ids_after - ids_before
+
+            if (
+                cupy.is_coreform()
+                and filter_sheet_bodies
+                and geometry == cupy.geometry.volume
+            ):
+                ids_new = {id for id in ids_new if not self.is_sheet_body(id)}
+
+            geometry_objects = self.get_items(geometry, item_ids=ids_new)
+            create_objects[geometry] = geometry_objects
+
+        return create_objects
+
     def display_in_cubit(self, labels=[], delay=0.5, testing=False):
         """Save the state to a cubit file and open cubit with that file.
         Additionally labels can be displayed in cubit to simplify the mesh

tests/test_cubitpy.py

@@ -2018,3 +2018,57 @@ def test_yaml_with_exo_export_fsi():
 
     # Compare the input file created for 4C.
     compare_yaml(cubit, mesh_in_exo=True)
+
+
+def test_cmd_return():
+    """Test the cmd_return function of CubitPy."""
+
+    cubit = CubitPy()
+
+    center = cubit.cmd_return("create vertex 0 0 0", cupy.geometry.vertex)
+    assert center.get_geometry_type() == cupy.geometry.vertex
+    assert center.id() == 1
+
+    arc_1 = cubit.cmd_return(
+        f"create curve arc center vertex {center.id()} radius 1 full",
+        cupy.geometry.curve,
+    )
+    assert arc_1.get_geometry_type() == cupy.geometry.curve
+    assert arc_1.id() == 1
+
+    center = cubit.cmd_return("create vertex 0.1 0 0", cupy.geometry.vertex)
+    assert center.get_geometry_type() == cupy.geometry.vertex
+    assert center.id() == 3
+
+    arc_2 = cubit.cmd_return(
+        f"create curve arc center vertex {center.id()} radius 2 full",
+        cupy.geometry.curve,
+    )
+    assert arc_2.get_geometry_type() == cupy.geometry.curve
+    assert arc_2.id() == 2
+
+    # We check the volume here as well, as in CoreForm a sheet body is created here that Cubit
+    # internally handles as a volume. But, we don't want this volume returned here.
+    create_surface_geometry = cubit.cmd_return_dict(
+        f"create surface curve {arc_1.id()} {arc_2.id()}",
+        [cupy.geometry.surface, cupy.geometry.volume],
+    )
+    for surface in create_surface_geometry[cupy.geometry.surface]:
+        assert surface.get_geometry_type() == cupy.geometry.surface
+    assert [item.id() for item in create_surface_geometry[cupy.geometry.surface]] == [1]
+    assert len(create_surface_geometry[cupy.geometry.volume]) == 0
+
+    sweep_geometry = cubit.cmd_return_dict(
+        f"sweep surface {surface.id()} perpendicular distance 2",
+        [
+            cupy.geometry.vertex,
+            cupy.geometry.curve,
+            cupy.geometry.surface,
+            cupy.geometry.volume,
+        ],
+    )
+    assert len(sweep_geometry) == 4
+    assert [item.id() for item in sweep_geometry[cupy.geometry.vertex]] == [5, 6]
+    assert [item.id() for item in sweep_geometry[cupy.geometry.curve]] == [3, 4, 5, 6]
+    assert [item.id() for item in sweep_geometry[cupy.geometry.surface]] == [2, 3, 4]
+    assert [item.id() for item in sweep_geometry[cupy.geometry.volume]] == [1]