diff --git a/Common/Constants/include/CommonConstants/make_pdg_header.py b/Common/Constants/include/CommonConstants/make_pdg_header.py index a94450e659acd..f83c44bb401db 100755 --- a/Common/Constants/include/CommonConstants/make_pdg_header.py +++ b/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 , 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()