Fix openmc-dev#2994, non-existent path causes segmentation fault when…

… saving plot (openmc-dev#3038) Co-authored-by: Patrick Shriwise <[email protected]> Co-authored-by: Paul Romano <[email protected]>
openmsr · Jul 18, 2024 · 5a10426 · 5a10426
1 parent bfbcc0b
commit 5a10426
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 1 deletion.
diff --git a/src/plot.cpp b/src/plot.cpp
@@ -315,6 +315,11 @@ void Plot::set_output_path(pugi::xml_node plot_node)
   } else {
     filename = fmt::format("plot_{}", id());
   }
+  const std::string dir_if_present =
+    filename.substr(0, filename.find_last_of("/") + 1);
+  if (dir_if_present.size() > 0 && !dir_exists(dir_if_present)) {
+    fatal_error(fmt::format("Directory '{}' does not exist!", dir_if_present));
+  }
   // add appropriate file extension to name
   switch (type_) {
   case PlotType::slice:

diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py
@@ -11,7 +11,7 @@ def myplot():
     plot.width = (100., 100.)
     plot.origin = (2., 3., -10.)
     plot.pixels = (500, 500)
-    plot.filename = 'myplot'
+    plot.filename = './not-a-dir/myplot'
     plot.type = 'slice'
     plot.basis = 'yz'
     plot.background = 'black'
@@ -221,3 +221,26 @@ def test_voxel_plot_roundtrip():
     assert new_plot.origin == plot.origin
     assert new_plot.width == plot.width
     assert new_plot.color_by == plot.color_by
+
+
+def test_plot_directory(run_in_tmpdir):
+    pwr_pin = openmc.examples.pwr_pin_cell()
+
+    # create a standard plot, expected to work
+    plot = openmc.Plot()
+    plot.filename = 'plot_1'
+    plot.type = 'slice'
+    plot.pixels = (10, 10)
+    plot.color_by = 'material'
+    plot.width = (100., 100.)
+    pwr_pin.plots = [plot]
+    pwr_pin.plot_geometry()
+
+    # use current directory, also expected to work
+    plot.filename = './plot_1'
+    pwr_pin.plot_geometry()
+
+    # use a non-existent directory, should raise an error
+    plot.filename = './not-a-dir/plot_1'
+    with pytest.raises(RuntimeError, match='does not exist'):
+        pwr_pin.plot_geometry()