more flexibility for mp

Alignment/OfflineValidation/python/TkAlAllInOneTool/alignment.py

@@ -1,7 +1,9 @@
-import configTemplates
-from helperFunctions import replaceByMap, parsecolor, parsestyle
+import collections
 import os
 import re
+
+import configTemplates
+from helperFunctions import replaceByMap, parsecolor, parsestyle
 from TkAlExceptions import AllInOneError
 
 class Alignment:
@@ -29,7 +31,7 @@ def __init__(self, name, config, runGeomComp = "1"):
             raise AllInOneError("section %s not found. Please define the "
                                   "alignment!"%section)
         config.checkInput(section,
-                          knownSimpleOptions = ['globaltag', 'style', 'color', 'title', 'mp', 'hp', 'sm'],
+                          knownSimpleOptions = ['globaltag', 'style', 'color', 'title', 'mp', 'mp_alignments', 'mp_deformations', 'hp', 'sm'],
                           knownKeywords = ['condition'])
         self.name = name
         if config.exists(section,"title"):
@@ -70,25 +72,71 @@ def __shorthandExists(self, theRcdName, theShorthand):
     def __getConditions( self, theConfig, theSection ):
         conditions = []
         for option in theConfig.options( theSection ):
-            if option == "mp":
+            if option in ("mp", "mp_alignments", "mp_deformations"):
                 condPars = theConfig.get(theSection, option).split(",")
-                if len(condPars) > 1:
-                    raise AllInOneError("Only one argument accepted for mp (should be the job number)")
-                number, = condPars
+                condPars = [_.strip() for _ in condPars]
+                if len(condPars) == 1:
+                    number, = condPars
+                    jobm = None
+                elif len(condPars) == 2:
+                    number, jobm = condPars
+                else:
+                    raise AllInOneError("Up to 2 arguments accepted for {} (job number, and optionally jobm index)".format(option))
+
+                if option == "mp":
+                    alignments = True
+                    deformations = True
+                elif option == "mp_alignments":
+                    alignments = True
+                    deformations = False
+                    option = "mp"
+                elif option == "mp_deformations":
+                    alignments = False
+                    deformations = True
+                    option = "mp"
+                else:
+                    assert False
+
                 folder = "/afs/cern.ch/cms/CAF/CMSALCA/ALCA_TRACKERALIGN/MP/MPproduction/{}{}/".format(option, number)
                 if not os.path.exists(folder):
                     raise AllInOneError(folder+" does not exist.")
-                dbfile = os.path.join(folder, "jobData/jobm/alignments_MP.db")
+                folder = os.path.join(folder, "jobData")
+                jobmfolders = set()
+                if jobm is None:
+                    for filename in os.listdir(folder):
+                        if re.match("jobm([0-9]*)", filename) and os.path.isdir(os.path.join(folder, filename)):
+                            jobmfolders.add(filename)
+                    if len(jobmfolders) == 0:
+                        raise AllInOneError("No jobm or jobm(number) folder in {}".format(folder))
+                    elif len(jobmfolders) == 1:
+                        folder = os.path.join(folder, jobmfolders.pop())
+                    else:
+                        raise AllInOneError(
+                                            "Multiple jobm or jobm(number) folders in {}\n".format(folder)
+                                            + ", ".join(jobmfolders) + "\n"
+                                            + "Please specify 0 for jobm, or a number for one of the others."
+                                           )
+                elif jobm == "0":
+                    folder = os.path.join(folder, "jobm")
+                    if os.path.exists(folder + "0"):
+                        raise AllInOneError("Not set up to handle a folder named jobm0")
+                else:
+                    folder = os.path.join(folder, "jobm{}".format(jobm))
+
+                dbfile = os.path.join(folder, "alignments_MP.db")
                 if not os.path.exists(dbfile):
-                    raise AllInOneError("No file {}.  Maybe your alignment folder is corrupted?".format(dbfile))
-                conditions.append({"rcdName": "TrackerAlignmentRcd",
-                                   "connectString": "sqlite_file:"+dbfile,
-                                   "tagName": "Alignments",
-                                   "labelName": ""})
-                conditions.append({"rcdName": "TrackerSurfaceDeformationRcd",
-                                   "connectString": "sqlite_file:"+dbfile,
-                                   "tagName": "Deformations",
-                                   "labelName": ""})
+                    raise AllInOneError("No file {}.  Maybe your alignment folder is corrupted, or maybe you specified the wrong jobm?".format(dbfile))
+
+                if alignments:
+                    conditions.append({"rcdName": "TrackerAlignmentRcd",
+                                       "connectString": "sqlite_file:"+dbfile,
+                                       "tagName": "Alignments",
+                                       "labelName": ""})
+                if deformations:
+                    conditions.append({"rcdName": "TrackerSurfaceDeformationRcd",
+                                       "connectString": "sqlite_file:"+dbfile,
+                                       "tagName": "Deformations",
+                                       "labelName": ""})
 
             elif option in ("hp", "sm"):
                 condPars = theConfig.get(theSection, option).split(",")
@@ -164,6 +212,13 @@ def __getConditions( self, theConfig, theSection ):
                                    "connectString": condPars[0].strip(),
                                    "tagName": condPars[1].strip(),
                                    "labelName": condPars[2].strip()})
+
+        rcdnames = collections.Counter(condition["rcdName"] for condition in conditions)
+        if max(rcdnames.values()) >= 2:
+            raise AllInOneError("Some conditions are specified multiple times (possibly through mp or hp options)!\n"
+                                + ", ".join(rcdname for rcdname, count in rcdnames.iteritems() if count >= 2))
+
+
         return conditions
 
     def __testDbExist(self, dbpath):