Resolve relative input paths relative to XML files (#4031)

Co-authored-by: GuySten <guyste@post.bgu.ac.il>

openmc/geometry.py

@@ -11,6 +11,7 @@ import openmc
 import openmc._xml as xml
 from .plots import add_plot_params
 from .checkvalue import check_type, check_less_than, check_greater_than, PathLike
+from .utility_funcs import set_xml_input_path
 
 
 class Geometry:
@@ -290,11 +291,12 @@ class Geometry:
         if isinstance(materials, (str, os.PathLike)):
             materials = openmc.Materials.from_xml(materials)
 
-        parser = ET.XMLParser(huge_tree=True)
-        tree = ET.parse(path, parser=parser)
-        root = tree.getroot()
+        with set_xml_input_path(path):
+            parser = ET.XMLParser(huge_tree=True)
+            tree = ET.parse(path, parser=parser)
+            root = tree.getroot()
 
-        return cls.from_xml_element(root, materials)
+            return cls.from_xml_element(root, materials)
 
     def find(self, point) -> list:
         """Find cells/universes/lattices which contain a given point

openmc/material.py

@@ -20,7 +20,7 @@ import openmc.data
 import openmc.checkvalue as cv
 from ._xml import clean_indentation, get_elem_list, get_text
 from .mixin import IDManagerMixin
-from .utility_funcs import input_path
+from .utility_funcs import input_path, set_xml_input_path
 from . import waste
 from openmc.checkvalue import PathLike
 from openmc.stats import Univariate, Discrete, Mixture, Tabular
@@ -2286,11 +2286,12 @@ class Materials(cv.CheckedList):
             Materials collection
 
         """
-        parser = ET.XMLParser(huge_tree=True)
-        tree = ET.parse(path, parser=parser)
-        root = tree.getroot()
+        with set_xml_input_path(path):
+            parser = ET.XMLParser(huge_tree=True)
+            tree = ET.parse(path, parser=parser)
+            root = tree.getroot()
 
-        return cls.from_xml_element(root)
+            return cls.from_xml_element(root)
 
 
     def deplete(

openmc/model/model.py

@@ -25,7 +25,7 @@ from openmc.checkvalue import (check_type, check_value, check_greater_than,
                                check_length, PathLike)
 from openmc.exceptions import InvalidIDError
 from openmc.plots import add_plot_params, _BASIS_INDICES, id_map_to_rgb
-from openmc.utility_funcs import change_directory
+from openmc.utility_funcs import change_directory, set_xml_input_path
 
 
 # Protocol for a function that is passed to search_keff
@@ -384,30 +384,31 @@ class Model:
         path : PathLike
             Path to model.xml file
         """
-        parser = ET.XMLParser(huge_tree=True)
-        tree = ET.parse(path, parser=parser)
-        root = tree.getroot()
+        with set_xml_input_path(path):
+            parser = ET.XMLParser(huge_tree=True)
+            tree = ET.parse(path, parser=parser)
+            root = tree.getroot()
 
-        model = cls()
+            model = cls()
 
-        desc_elem = root.find('description')
-        if desc_elem is not None and desc_elem.text:
-            model.description = desc_elem.text
+            desc_elem = root.find('description')
+            if desc_elem is not None and desc_elem.text:
+                model.description = desc_elem.text
 
-        meshes = {}
-        model.settings = openmc.Settings.from_xml_element(
-            root.find('settings'), meshes)
-        model.materials = openmc.Materials.from_xml_element(
-            root.find('materials'))
-        model.geometry = openmc.Geometry.from_xml_element(
-            root.find('geometry'), model.materials)
+            meshes = {}
+            model.settings = openmc.Settings.from_xml_element(
+                root.find('settings'), meshes)
+            model.materials = openmc.Materials.from_xml_element(
+                root.find('materials'))
+            model.geometry = openmc.Geometry.from_xml_element(
+                root.find('geometry'), model.materials)
 
-        if root.find('tallies') is not None:
-            model.tallies = openmc.Tallies.from_xml_element(
-                root.find('tallies'), meshes)
+            if root.find('tallies') is not None:
+                model.tallies = openmc.Tallies.from_xml_element(
+                    root.find('tallies'), meshes)
 
-        if root.find('plots') is not None:
-            model.plots = openmc.Plots.from_xml_element(root.find('plots'))
+            if root.find('plots') is not None:
+                model.plots = openmc.Plots.from_xml_element(root.find('plots'))
 
         return model
 

openmc/plots.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from collections.abc import Iterable, Mapping, Sequence
 from numbers import Integral, Real
 from pathlib import Path
@@ -14,6 +15,7 @@ from openmc.checkvalue import PathLike
 
 from ._xml import clean_indentation, get_elem_list, get_text
 from .mixin import IDManagerMixin
+from .utility_funcs import set_xml_input_path
 
 _BASES = {'xy', 'xz', 'yz'}
 
@@ -2190,7 +2192,7 @@ class Plots(cv.CheckedList):
         tree.write(str(p), xml_declaration=True, encoding='utf-8')
 
     @classmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem) -> Plots:
         """Generate plots collection from XML file
 
         Parameters
@@ -2224,7 +2226,7 @@ class Plots(cv.CheckedList):
         return plots
 
     @classmethod
-    def from_xml(cls, path='plots.xml'):
+    def from_xml(cls, path: PathLike = 'plots.xml') -> Plots:
         """Generate plots collection from XML file
 
         Parameters
@@ -2238,7 +2240,8 @@ class Plots(cv.CheckedList):
             Plots collection
 
         """
-        parser = ET.XMLParser(huge_tree=True)
-        tree = ET.parse(path, parser=parser)
-        root = tree.getroot()
-        return cls.from_xml_element(root)
+        with set_xml_input_path(path):
+            parser = ET.XMLParser(huge_tree=True)
+            tree = ET.parse(path, parser=parser)
+            root = tree.getroot()
+            return cls.from_xml_element(root)

openmc/settings.py

@@ -15,7 +15,7 @@ from openmc.stats.multivariate import MeshSpatial
 from ._xml import clean_indentation, get_elem_list, get_text
 from .mesh import _read_meshes, RegularMesh, MeshBase
 from .source import SourceBase, MeshSource, IndependentSource
-from .utility_funcs import input_path
+from .utility_funcs import input_path, set_xml_input_path
 from .volume import VolumeCalculation
 from .weight_windows import WeightWindows, WeightWindowGenerator, WeightWindowsList
 
@@ -2775,8 +2775,9 @@ class Settings:
             Settings object
 
         """
-        parser = ET.XMLParser(huge_tree=True)
-        tree = ET.parse(path, parser=parser)
-        root = tree.getroot()
-        meshes = _read_meshes(root)
-        return cls.from_xml_element(root, meshes)
+        with set_xml_input_path(path):
+            parser = ET.XMLParser(huge_tree=True)
+            tree = ET.parse(path, parser=parser)
+            root = tree.getroot()
+            meshes = _read_meshes(root)
+            return cls.from_xml_element(root, meshes)

openmc/tallies.py

@@ -36,6 +36,7 @@ from openmc.arithmetic import (
 from ._sparse_compat import lil_array
 from ._xml import clean_indentation, get_elem_list, get_text
 from .mixin import IDManagerMixin
+from .utility_funcs import set_xml_input_path
 from .mesh import MeshBase
 
 
@@ -3938,7 +3939,8 @@ class Tallies(cv.CheckedList):
             Tallies object
 
         """
-        parser = ET.XMLParser(huge_tree=True)
-        tree = ET.parse(path, parser=parser)
-        root = tree.getroot()
-        return cls.from_xml_element(root)
+        with set_xml_input_path(path):
+            parser = ET.XMLParser(huge_tree=True)
+            tree = ET.parse(path, parser=parser)
+            root = tree.getroot()
+            return cls.from_xml_element(root)

openmc/utility_funcs.py

@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+from contextvars import ContextVar
 import os
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -9,6 +10,19 @@ import openmc
 from .checkvalue import PathLike
 
 
+_XML_INPUT_PATH = ContextVar('_XML_INPUT_PATH', default=None)
+
+
+@contextmanager
+def set_xml_input_path(path: PathLike):
+    """Set the containing XML directory for resolving input paths."""
+    token = _XML_INPUT_PATH.set(Path(path).resolve().parent)
+    try:
+        yield
+    finally:
+        _XML_INPUT_PATH.reset(token)
+
+
 @contextmanager
 def change_directory(working_dir: PathLike | None = None, *, tmpdir: bool = False):
     """Context manager for executing in a provided working directory
@@ -56,7 +70,11 @@ def input_path(filename: PathLike) -> Path:
 
     """
     if openmc.config['resolve_paths']:
-        return Path(filename).resolve()
+        path = Path(filename)
+        xml_dir = _XML_INPUT_PATH.get()
+        if xml_dir is not None and not path.is_absolute():
+            path = xml_dir / path
+        return path.resolve()
     else:
         return Path(filename)
 

tests/unit_tests/test_model.py

@@ -219,6 +219,25 @@ def test_from_xml(run_in_tmpdir, pin_model_attributes):
     assert test_model.is_initialized is False
 
 
+@pytest.mark.parametrize('resolve_paths', [True, False])
+def test_xml_input_paths(tmp_path, monkeypatch, resolve_paths):
+    """Input paths are interpreted relative to the containing XML file."""
+    with openmc.config.patch('resolve_paths', False):
+        univ = openmc.DAGMCUniverse('dagmc.h5m')
+        model = openmc.Model(geometry=openmc.Geometry(univ))
+        model_dir = tmp_path / 'model'
+        model.export_to_model_xml(model_dir)
+
+    monkeypatch.chdir(tmp_path)
+    with openmc.config.patch('resolve_paths', resolve_paths):
+        openmc.reset_auto_ids()
+        loaded_model = openmc.Model.from_model_xml('model/model.xml')
+        expected = (model_dir / 'dagmc.h5m').resolve()
+        if not resolve_paths:
+            expected = Path('dagmc.h5m')
+        assert loaded_model.geometry.root_universe.filename == expected
+
+
 def test_init_finalize_lib(run_in_tmpdir, pin_model_attributes, mpi_intracomm):
     # We are going to init and then make sure data is loaded
     mats, geom, settings, tals, plots, _, _ = pin_model_attributes