From 18c8db4088c878db5bf1748f30476a6a2f0128e3 Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Wed, 17 Apr 2024 12:31:05 +0100
Subject: [PATCH 1/3] returning number of volumes

---
 src/cad_to_dagmc/core.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/cad_to_dagmc/core.py b/src/cad_to_dagmc/core.py
index 6bf8ebb..0a59d86 100644
--- a/src/cad_to_dagmc/core.py
+++ b/src/cad_to_dagmc/core.py
@@ -396,7 +396,7 @@ def add_stp_file(
         filename: str,
         scale_factor: float = 1.0,
         material_tags: typing.Optional[typing.Iterable[str]] = None,
-    ):
+    ) -> int:
         """Loads the parts from stp file into the model.
 
         Args:
@@ -410,6 +410,9 @@ def add_stp_file(
                 used to increase the size or decrease the size of the geometry.
                 Useful when converting the geometry to cm for use in neutronics
                 simulations.
+
+        Returns:
+            int: number of volumes in the stp file.
         """
         part = importers.importStep(str(filename)).val()
 
@@ -417,7 +420,10 @@ def add_stp_file(
             scaled_part = part
         else:
             scaled_part = part.scale(scale_factor)
-        self.add_cadquery_object(cadquery_object=scaled_part, material_tags=material_tags)
+        return self.add_cadquery_object(
+            cadquery_object=scaled_part,
+            material_tags=material_tags
+        )
 
     def add_cadquery_object(
         self,
@@ -425,7 +431,7 @@ def add_cadquery_object(
             cq.assembly.Assembly, cq.occ_impl.shapes.Compound, cq.occ_impl.shapes.Solid
         ],
         material_tags: typing.Optional[typing.Iterable[str]] = None,
-    ):
+    ) -> int:
         """Loads the parts from CadQuery object into the model.
 
         Args:
@@ -436,6 +442,9 @@ def add_cadquery_object(
                 same order as the volumes in the geometry added (STP file and
                 CadQuery objects) and match the material tags used in the
                 neutronics code (e.g. OpenMC).
+
+        Returns:
+            int: number of volumes in the stp file.
         """
 
         if isinstance(cadquery_object, cq.assembly.Assembly):
@@ -451,6 +460,8 @@ def add_cadquery_object(
             self.material_tags = self.material_tags + material_tags
         self.parts = self.parts + iterable_solids
 
+        return len(iterable_solids)
+
     def export_unstructured_mesh_file(
         self,
         filename: str = "umesh.h5m",

From dff023c94ae3d325eaed32f9805e998ead46feb7 Mon Sep 17 00:00:00 2001
From: shimwell <shimwell@users.noreply.github.com>
Date: Wed, 17 Apr 2024 11:31:44 +0000
Subject: [PATCH 2/3] [skip ci] Apply formatting changes

---
 src/cad_to_dagmc/core.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/cad_to_dagmc/core.py b/src/cad_to_dagmc/core.py
index 0a59d86..d0408bd 100644
--- a/src/cad_to_dagmc/core.py
+++ b/src/cad_to_dagmc/core.py
@@ -420,10 +420,7 @@ def add_stp_file(
             scaled_part = part
         else:
             scaled_part = part.scale(scale_factor)
-        return self.add_cadquery_object(
-            cadquery_object=scaled_part,
-            material_tags=material_tags
-        )
+        return self.add_cadquery_object(cadquery_object=scaled_part, material_tags=material_tags)
 
     def add_cadquery_object(
         self,

From 94d9f537f8d77bebf35bc7b213737b84f4e9bfcf Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Wed, 17 Apr 2024 15:34:34 +0100
Subject: [PATCH 3/3] added tests for returned number of solids

---
 tests/test_python_api.py | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/tests/test_python_api.py b/tests/test_python_api.py
index 78abc0c..50facf1 100644
--- a/tests/test_python_api.py
+++ b/tests/test_python_api.py
@@ -106,3 +106,35 @@ def test_h5m_file_tags():
         2: "mat:mat2",
         3: "mat:mat3",
     }
+
+
+def test_add_cadquery_object_returned_volumes():
+    """Checks that a add_cadquery_object method returns the correct number of volumes"""
+
+    sphere1 = cq.Workplane().sphere(20)
+    sphere2 = cq.Workplane().moveTo(100, 100).sphere(20)
+    sphere3 = cq.Workplane().moveTo(-100, -100).sphere(20)
+
+    c2d = CadToDagmc()
+    vols = c2d.add_cadquery_object(sphere1)
+    assert vols == 1
+
+    assembly = cq.Assembly()
+    assembly.add(sphere1)
+    assembly.add(sphere2)
+    assembly.add(sphere3)
+    c2d = CadToDagmc()
+    vols = c2d.add_cadquery_object(assembly)
+    assert vols == 3
+
+
+def test_add_stp_file_returned_volumes():
+    """Checks that a add_stp_file method returns the correct number of volumes"""
+
+    c2d = CadToDagmc()
+    vols = c2d.add_stp_file("tests/curved_extrude.stp")
+    assert vols == 1
+
+    c2d = CadToDagmc()
+    vols = c2d.add_stp_file("tests/two_disconnected_cubes.stp")
+    assert vols == 2