diff --git a/src/cad_to_dagmc/core.py b/src/cad_to_dagmc/core.py index 6bf8ebb..d0408bd 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,7 @@ 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 +428,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 +439,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 +457,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", 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