Add main

Author: vkucera

Common/Constants/include/CommonConstants/make_pdg_header.py

@@ -23,7 +23,6 @@
 
 import ROOT  # pylint: disable=import-error
 
-name_script = os.path.basename(__file__)
 
 # Enum of PDG_t particles
 class PdgROOT(Enum):
@@ -149,6 +148,7 @@ class Pdg(Enum):
     kHyperHelium4Sigma = 1110020040
     kLambda1520_Py = 102134  # PYTHIA code different from PDG
 
+
 dbPdg = ROOT.o2.O2DatabasePDG
 
 
@@ -159,49 +159,58 @@ def mass(code):
     return dbPdg.Mass(code, success)
 
 
-def declare_mass(pdg, type="double") -> str:
+def declare_mass(pdg, val_type="double") -> str:
     """Returns a C++ declaration of a particle mass constant."""
-    return f"constexpr {type} Mass{pdg.name[1:]} = {mass(pdg.value)};\n"
-
-
-# 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
-"""
-
-# 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)
+    return f"constexpr {val_type} Mass{pdg.name[1:]} = {mass(pdg.value)};\n"
+
+
+def main():
+    """Main function"""
+
+    name_script = os.path.basename(__file__)
+
+    # Comment at the beginning of the output
+    str_block_begin = "// BEGINNING OF THE GENERATED BLOCK."
+    # Preamble with instructions
+    str_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.\n"
+    )
+    # 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"
+        "\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
+    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"
+
+    # 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_block_preamble, str_enum, str_mass_o2, str_mass_root, str_block_end))
+    print(str_header)
+
+
+if __name__ == "__main__":
+    main()