tests/unit_tests/test_geo.py

0001 """TGeo walking in pyrobird.cern_root and the `pyrobird geo` commands.
0002
0003 The fake ROOT module below mimics PyROOT where it matters: TGeoIterator.Next()
0004 ends with a null proxy that is falsy but not None, and removing a node while
0005 the iterator runs is an error (ROOT crashes).
0006 """
0007 import os
0008 import subprocess
0009 import sys
0010 import types
0011
0012 import pytest
0013 from click.testing import CliRunner
0014
0015
0016 class NullProxy:
0017     """PyROOT's stand-in for a null TGeoNode*: falsy, and not None."""
0018
0019     def __bool__(self):
0020         return False
0021
0022
0023 class FakeVolume:
0024     def __init__(self, name, daughters=()):
0025         self.name = name
0026         self.daughters = list(daughters)
0027         self.color = None
0028
0029     def GetName(self):
0030         return self.name
0031
0032     def SetLineColor(self, color):
0033         self.color = color
0034
0035     def RemoveNode(self, node):
0036         if FakeIterator.active:
0037             raise RuntimeError("RemoveNode while TGeoIterator runs (ROOT crashes here)")
0038         self.daughters.remove(node)
0039
0040
0041 class FakeNode:
0042     def __init__(self, name, volume, mother):
0043         self.name = name
0044         self.volume = volume
0045         self.mother = mother
0046
0047     def GetVolume(self):
0048         return self.volume
0049
0050     def GetMotherVolume(self):
0051         return self.mother
0052
0053
0054 class FakeTString:
0055     def __init__(self):
0056         self.value = ""
0057
0058     def __str__(self):
0059         return self.value
0060
0061
0062 class FakeIterator:
0063     active = False
0064
0065     def __init__(self, top_volume):
0066         # Depth-first list of (path, level, node), honoring Skip()
0067         self.top = top_volume
0068         self.stack = [(f"/{top_volume.name}_1/{n.name}", 1, n) for n in reversed(top_volume.daughters)]
0069         self.current = None
0070         self.skip_next = False
0071         self.skips = 0
0072         FakeIterator.active = True
0073
0074     def Next(self):
0075         if self.current is not None and not self.skip_next:
0076             path, level, node = self.current
0077             self.stack.extend((f"{path}/{d.name}", level + 1, d) for d in reversed(node.volume.daughters))
0078         self.skip_next = False
0079         if not self.stack:
0080             FakeIterator.active = False
0081             self.current = None
0082             return NullProxy()
0083         self.current = self.stack.pop()
0084         return self.current[2]
0085
0086     def GetPath(self, tstring):
0087         tstring.value = self.current[0]
0088
0089     def GetLevel(self):
0090         return self.current[1]
0091
0092     def Skip(self):
0093         self.skip_next = True
0094         self.skips += 1
0095
0096
0097 class FakeGeoManager:
0098     def __init__(self, top):
0099         self.top = top
0100         self.exported = None
0101
0102     def GetMasterVolume(self):
0103         return self.top
0104
0105     def GetTopVolume(self):
0106         return self.top
0107
0108     def GetNNodes(self):
0109         return 4
0110
0111     def CleanGarbage(self):
0112         pass
0113
0114     def Export(self, file_name):
0115         self.exported = file_name
0116
0117
0118 def make_geometry():
0119     """TOP -> DIRC -> DIRCModule_0, DIRCModule_1; TOP -> Other"""
0120     top = FakeVolume("TOP")
0121     dirc = FakeVolume("DIRC")
0122     other = FakeVolume("Other")
0123     module0 = FakeVolume("DIRCModule_0")
0124     module1 = FakeVolume("DIRCModule_1")
0125     dirc_node = FakeNode("DIRC_1", dirc, top)
0126     other_node = FakeNode("Other_1", other, top)
0127     dirc.daughters = [FakeNode("DIRCModule_0_1", module0, dirc), FakeNode("DIRCModule_1_1", module1, dirc)]
0128     top.daughters = [dirc_node, other_node]
0129     return FakeGeoManager(top)
0130
0131
0132 @pytest.fixture
0133 def fake_root(monkeypatch):
0134     geo_manager = make_geometry()
0135     module = types.ModuleType("ROOT")
0136     module.gErrorIgnoreLevel = 0
0137     module.kFatal = 6000
0138     module.kMagenta = 616
0139     module.TGeoManager = types.SimpleNamespace(Import=lambda file_name: geo_manager)
0140     module.TGeoIterator = FakeIterator
0141     module.TString = FakeTString
0142     # Replace only the ROOT entry: modules imported meanwhile must stay loaded
0143     monkeypatch.setitem(sys.modules, "ROOT", module)
0144     yield geo_manager
0145     FakeIterator.active = False
0146
0147
0148 def test_process_file_walk_ends_on_null_proxy(fake_root):
0149     from pyrobird.cern_root import tgeo_process_file
0150     patterns = ["*/DIRCModule_*", "*/Other*"]
0151
0152     removed = tgeo_process_file("in.root", "out.root", patterns)
0153
0154     # Each pattern removes the first node it matches; removal waits for the walk to end
0155     assert removed == ["/TOP_1/DIRC_1/DIRCModule_0_1", "/TOP_1/Other_1"]
0156     top = fake_root.top
0157     assert [n.name for n in top.daughters] == ["DIRC_1"]
0158     assert [n.name for n in top.daughters[0].volume.daughters] == ["DIRCModule_1_1"]
0159     assert fake_root.exported == "out.root"
0160     # The caller's list is not modified
0161     assert patterns == ["*/DIRCModule_*", "*/Other*"]
0162
0163
0164 def test_info_prints_paths_to_depth(fake_root):
0165     from pyrobird.cern_root import tgeo_info
0166     lines = []
0167
0168     printed = tgeo_info("in.root", max_depth=1, echo=lines.append)
0169
0170     assert printed == 2
0171     assert lines == ["Top volume: TOP", "Nodes: 4", "1 /TOP_1/DIRC_1", "1 /TOP_1/Other_1"]
0172     # info never edits the geometry
0173     assert fake_root.exported is None
0174     assert len(fake_root.top.daughters) == 2
0175
0176
0177 def test_info_command_does_not_write_files(fake_root, tmp_path, monkeypatch):
0178     from pyrobird.cli.geo import geo
0179     monkeypatch.chdir(tmp_path)
0180
0181     result = CliRunner().invoke(geo, ["info", "detector.root", "--max-depth", "2"])
0182
0183     assert result.exit_code == 0, result.output
0184     assert "2 /TOP_1/DIRC_1/DIRCModule_0_1" in result.output
0185     assert list(tmp_path.iterdir()) == []
0186
0187
0188 def test_process_command_help():
0189     from pyrobird.cli.geo import geo
0190     result = CliRunner().invoke(geo, ["process", "--help"])
0191     assert result.exit_code == 0
0192     assert "Remove nodes from a TGeo geometry file" in result.output
0193     assert "YAML rules file" in result.output
0194
0195
0196 # --- Real PyROOT, when the interpreter has it ---------------------------------
0197
0198 MAKE_GEOMETRY = """
0199 import ROOT
0200 ROOT.gErrorIgnoreLevel = ROOT.kFatal
0201 gm = ROOT.TGeoManager("g", "g")
0202 mat = ROOT.TGeoMaterial("Vac", 0, 0, 0)
0203 med = ROOT.TGeoMedium("Vac", 1, mat)
0204 top = gm.MakeBox("TOP", med, 100, 100, 100)
0205 gm.SetTopVolume(top)
0206 dirc = gm.MakeBox("DIRC", med, 10, 10, 10)
0207 module = gm.MakeBox("DIRCModule_0", med, 1, 1, 1)
0208 dirc.AddNode(module, 1)
0209 top.AddNode(dirc, 1)
0210 other = gm.MakeBox("Other", med, 1, 1, 1)
0211 top.AddNode(other, 1)
0212 gm.CloseGeometry()
0213 gm.Export(__import__("sys").argv[1])
0214 # A TGeoManager built in Python crashes ROOT at interpreter teardown
0215 __import__("os")._exit(0)
0216 """
0217
0218 LIST_PATHS = """
0219 import sys
0220 import ROOT
0221 from pyrobird.cern_root import tgeo_walk
0222 ROOT.gErrorIgnoreLevel = ROOT.kFatal
0223 gm = ROOT.TGeoManager.Import(sys.argv[1])
0224 print("\\n".join(path for path, _, _ in tgeo_walk(gm)))
0225 """
0226
0227
0228 def _run_python(code, *args):
0229     # TGeoManager is a process-wide singleton: give each step its own process
0230     env = dict(os.environ, PYTHONPATH=os.pathsep.join(sys.path))
0231     return subprocess.run([sys.executable, "-c", code, *args], capture_output=True, text=True,
0232                           timeout=120, env=env)
0233
0234
0235 def test_process_file_with_pyroot(tmp_path):
0236     pytest.importorskip("ROOT")
0237     source = str(tmp_path / "geo.root")
0238     output = str(tmp_path / "geo.edit.root")
0239     assert _run_python(MAKE_GEOMETRY, source).returncode == 0
0240
0241     code = ("import sys; from pyrobird.cern_root import tgeo_process_file; "
0242             "print(tgeo_process_file(sys.argv[1], sys.argv[2], ['*/DIRCModule_0*']))")
0243     result = _run_python(code, source, output)
0244     assert result.returncode == 0, result.stderr
0245     assert "DIRCModule_0" in result.stdout
0246
0247     listing = _run_python(LIST_PATHS, output)
0248     assert listing.returncode == 0, listing.stderr
0249     paths = listing.stdout.split()
0250     assert "TOP/DIRC_1" in paths
0251     assert "TOP/Other_1" in paths
0252     assert not any("DIRCModule_0" in path for path in paths)
0253
0254
0255 def test_info_depth_with_pyroot(tmp_path):
0256     pytest.importorskip("ROOT")
0257     source = str(tmp_path / "geo.root")
0258     assert _run_python(MAKE_GEOMETRY, source).returncode == 0
0259
0260     code = ("import sys; from pyrobird.cern_root import tgeo_info; "
0261             "print('printed', tgeo_info(sys.argv[1], max_depth=int(sys.argv[2])))")
0262     shallow = _run_python(code, source, "1")
0263     assert shallow.returncode == 0, shallow.stderr
0264     assert "printed 2" in shallow.stdout
0265     assert "DIRCModule_0" not in shallow.stdout
0266
0267     deep = _run_python(code, source, "2")
0268     assert "printed 3" in deep.stdout
0269     assert "2 TOP/DIRC_1/DIRCModule_0_1" in deep.stdout