Merge branch 'aiwieder/tuples_exist' into 'master'

Check if at least one DecayTree and one MCDecayTree are present in the output tuple

Closes LbAnalysisProductions#32

See merge request !33

src/LbAPCommon/options_parsing/__init__.py

@@ -13,31 +13,59 @@ import json
 
 import uproot
 
+from LbAPCommon.parsing import is_simulation_job
 
-def validate_options(json_file: str, ntuple_fn: str):
-    "Comparing the expected tuples with the output"
+
+def validate_options(json_file: str, ntuple_fn: str, job_data: dict, prod_data: dict):
+    """
+    Validate existence of expected TTrees and ensure at least one
+    DecayTree and MCDecayTree (for MC samples) exists in output files
+
+    :param json_file: json_file listing the expected TTrees
+    :param ntuple_fn: Local test output tuple to validate against
+    :param job_data: Job specific setup information from yaml parsing
+    :param prod_data: Entire production information from yaml parsing
+    :returns: Errors if any expected TTrees are not found in job output,
+              warnings if at least one (MC)DecayTree isn't found in job output
+    """
 
     with open(json_file, "rb") as fp:
         json_dump = json.load(fp)
 
     file = uproot.open(ntuple_fn)
-    tupleKeys = file.keys(cycle=False)
+    tuple_keys = file.keys(cycle=False)
+
+    errors = []
+    warnings = []
+
+    mc_decay_tree = any({"MCDecayTree" in key for key in tuple_keys})
+    decay_tree = any(
+        {("DecayTree" in key and "MCDecayTree" not in key) for key in tuple_keys}
+    )
+
+    if not decay_tree:
+        warnings.append(
+            "No DecayTree detected in the output file! Is this intentional?"
+        )
 
-    errorList = []
+    if is_simulation_job(prod_data, job_data["name"]) and not mc_decay_tree:
+        warnings.append(
+            "No MCDecayTree detected in the output file! Is this intentional?"
+        )
 
     # Checking DecayTreeTuples
-    for nTuple in json_dump["DecayTreeTuple"]:
-        if nTuple not in tupleKeys:
-            errorList.append(f"ERROR: DecayTreeTuple {nTuple} missing in test result")
+    for ntuple in json_dump["DecayTreeTuple"]:
+        if ntuple not in tuple_keys:
+            errors.append(f"ERROR: DecayTreeTuple {ntuple} missing in test result")
 
     # Checking MCDecayTreeTuples
-    for nTuple in json_dump["MCDecayTreeTuple"]:
-        if nTuple not in tupleKeys:
-            errorList.append(f"ERROR: MCDecayTreeTuple {nTuple} missing in test result")
+    for ntuple in json_dump["MCDecayTreeTuple"]:
+        if ntuple not in tuple_keys:
+            errors.append(f"ERROR: MCDecayTreeTuple {ntuple} missing in test result")
 
     # Checking MCDecayTreeTuples
-    for nTuple in json_dump["EventTuple"]:
-        if nTuple not in tupleKeys:
-            errorList.append(f"ERROR: EventTuple {nTuple} missing in test result")
+    for ntuple in json_dump["EventTuple"]:
+        if ntuple not in tuple_keys:
+            errors.append(f"ERROR: EventTuple {ntuple} missing in test result")
 
-    return errorList
+    return errors, warnings

src/LbAPCommon/parsing.py

@@ -243,6 +243,30 @@ def create_proc_pass_map(job_names, version, default_proc_pass="default"):
     return proc_pass_map
 
 
+def is_simulation_job(prod_data: dict, job_name: str):
+    """Determine if a job is using MC input or not.
+
+    :param prod_data: Entire production information from yaml parsing
+    :param job_name: Name of the job to determine if it's using MC input or not
+    :returns: True if the job is using MC input, False if it is not
+    """
+
+    job_dict = prod_data[job_name]
+    if "simulation" not in job_dict:
+        if "bk_query" in job_dict["input"]:
+            if "/mc/" in job_dict["input"]["bk_query"].lower():
+                return True
+            else:
+                return False
+        elif "job_name" in job_dict["input"]:
+            dependent_job = prod_data[job_name]["input"]["job_name"]
+            return is_simulation_job(prod_data, dependent_job)
+        else:
+            raise NotImplementedError(
+                "Input requires either a bookkeeping location or a previous job name"
+            )
+
+
 def parse_yaml(rendered_yaml):
     data1 = load(
         rendered_yaml, schema=MapPattern(Regex(RE_JOB_NAME), Any(), minimum_keys=1)