Improve the PDG header script

Common/Constants/include/CommonConstants/make_pdg_header.py

@@ -12,7 +12,7 @@
 # or submit itself to any jurisdiction.
 
 """!
-@brief  Generates the body of a C++ header with PDG codes and particle masses.
+@brief  Generates and updates the body of a C++ header with PDG codes and particle masses.
 @author Vít Kučera <vit.kucera@cern.ch>, Inha University
 @date   2023-09-21
 """
@@ -21,9 +21,12 @@
 from ctypes import c_bool
 from enum import Enum
 
-import ROOT  # pylint: disable=import-error
+try:
+    import ROOT  # pylint: disable=import-error
+    from ROOT import o2
+except (ModuleNotFoundError, ImportError) as exc:
+    raise OSError("O2 environment is not loaded.") from exc
 
-name_script = os.path.basename(__file__)
 
 # Enum of PDG_t particles
 class PdgROOT(Enum):
@@ -149,7 +152,8 @@ class Pdg(Enum):
     kHyperHelium4Sigma = 1110020040
     kLambda1520_Py = 102134  # PYTHIA code different from PDG
 
-dbPdg = ROOT.o2.O2DatabasePDG
+
+dbPdg = o2.O2DatabasePDG
 
 
 def mass(code):
@@ -159,49 +163,109 @@ def mass(code):
     return dbPdg.Mass(code, success)
 
 
-def declare_mass(pdg, type="double") -> str:
+def declare_mass(pdg, mass_type="double") -> str:
     """Returns a C++ declaration of a particle mass constant."""
-    return f"constexpr {type} Mass{pdg.name[1:]} = {mass(pdg.value)};\n"
+    return f"constexpr {mass_type} Mass{pdg.name[1:]} = {mass(pdg.value)};"
 
 
-# Comment at the beginning of the output
-str_block_begin = f"""// BEGINNING OF THE GENERATED BLOCK.
-// DO NOT EDIT THIS BLOCK DIRECTLY!
-// It has been generated by the {name_script} script.
-// For modifications, edit the script and generate this block again.
-"""
-# Comment at the end of the output
-str_block_end = """// END OF THE GENERATED BLOCK
-"""
-# Start of enum declarations of additional particles
-str_enum_head = """/// \\brief Declarations of named PDG codes of particles missing in ROOT PDG_t
-/// \\note Follow kCamelCase naming convention
-/// \\link https://root.cern/doc/master/TPDGCode_8h.html
-enum Pdg {
-"""
-# End of enum declarations of additional particles
-str_enum_foot = "};\n"
-# Documentation string for mass declarations of additional particles
-str_mass_o2_head = """/// \\brief Declarations of masses for additional particles
-"""
-# Documentation string for mass declarations of PDG_t particles
-str_mass_root_head = """/// \\brief Declarations of masses for particles in ROOT PDG_t
-"""
+def main():
+    """Main function"""
+
+    path_header = "PhysicsConstants.h"
+    name_script = os.path.basename(__file__)
+
+    # Comment at the beginning of the output
+    block_begin = "// BEGINNING OF THE GENERATED BLOCK."
+    # Comment at the end of the output
+    block_end = "// END OF THE GENERATED BLOCK"
+    # Preamble with instructions
+    block_preamble = (
+        "// DO NOT EDIT THIS BLOCK DIRECTLY!"
+        f"\n// It has been generated by the {name_script} script."
+        "\n// For modifications, edit the script and generate this block again."
+    )
+    # Start of enum declarations of additional particles
+    enum_o2_head = (
+        "/// \\brief Declarations of named PDG codes of particles missing in ROOT PDG_t"
+        "\n/// \\note Follow kCamelCase naming convention"
+        "\n/// \\link https://root.cern/doc/master/TPDGCode_8h.html"
+        "\nenum Pdg {"
+    )
+    # End of enum declarations of additional particles
+    enum_o2_foot = "};"
+    # Documentation string for mass declarations of additional particles
+    mass_o2_head = "/// \\brief Declarations of masses for additional particles"
+    # Documentation string for mass declarations of PDG_t particles
+    mass_root_head = "/// \\brief Declarations of masses for particles in ROOT PDG_t"
+
+    # Get header content before and after the generated block.
+    print(f'File "{path_header}" will be updated.')
+    try:
+        with open(path_header, encoding="utf-8") as file:
+            content_old = file.readlines()
+    except OSError as exc:
+        raise OSError(f'Failed to open file "{path_header}".') from exc
+    lines_header_before: list[str] = []
+    lines_header_after: list[str] = []
+    got_block_begin = False
+    got_block_end = False
+    for line in content_old:
+        line = line.strip()
+        if line == block_begin:
+            got_block_begin = True
+        if not got_block_begin:
+            lines_header_before.append(line)
+        if got_block_end:
+            lines_header_after.append(line)
+        if line == block_end:
+            got_block_end = True
+    if not got_block_begin:
+        raise ValueError("Did not find the beginning of the block.")
+    if not got_block_end:
+        raise ValueError("Did not find the end of the block.")
+
+    # Additional particles
+    lines_enum_o2: list[str] = [enum_o2_head]
+    lines_mass_o2: list[str] = [mass_o2_head]
+    for pdg_o2 in Pdg:
+        lines_enum_o2.append(f"  {pdg_o2.name} = {pdg_o2.value},")
+        lines_mass_o2.append(declare_mass(pdg_o2))
+    lines_enum_o2[-1] = lines_enum_o2[-1][:-1]  # Remove the last comma.
+    lines_enum_o2.append(enum_o2_foot)
+
+    # PDG_t particles
+    lines_mass_root: list[str] = [mass_root_head]
+    for pdg_root in PdgROOT:
+        lines_mass_root.append(declare_mass(pdg_root))
+
+    # Header body
+    content_new = "\n".join(
+        (
+            *lines_header_before,
+            block_begin,
+            block_preamble,
+            "",
+            *lines_enum_o2,
+            "",
+            *lines_mass_o2,
+            "",
+            *lines_mass_root,
+            "",
+            block_end,
+            *lines_header_after,
+            "",
+        )
+    )
+    # print(content_new)
+
+    # Overwrite the input file.
+    try:
+        with open(path_header, "w", encoding="utf-8") as file:
+            file.write(content_new)
+            print(f'File "{path_header}" has been overwritten.')
+    except OSError as exc:
+        raise OSError(f'Failed to write to file "{path_header}".') from exc
+
 
-# Additional particles
-str_enum = str_enum_head
-str_mass_o2 = str_mass_o2_head
-for c in Pdg:
-    str_enum += f"  {c.name} = {c.value},\n"
-    str_mass_o2 += declare_mass(c)
-str_enum = str_enum[:-2] + "\n"  # Remove the last comma.
-str_enum += str_enum_foot
-
-# PDG_t particles
-str_mass_root = str_mass_root_head
-for d in PdgROOT:
-    str_mass_root += declare_mass(d)
-
-# Header body
-str_header = "\n".join((str_block_begin, str_enum, str_mass_o2, str_mass_root, str_block_end))
-print(str_header)
+if __name__ == "__main__":
+    main()