provide functions for sprucing

PyConf/python/PyConf/reading.py

@@ -8,7 +8,7 @@
 # granted to it by virtue of its status as an Intergovernmental Organization  #
 # or submit itself to any jurisdiction.                                       #
 ###############################################################################
-from PyConf.packing import reco_locations, unpackers_map
+from PyConf.packing import reco_locations, unpackers_map, pp2mcp_locations
 from PyConf.components import force_location
 from PyConf.dataflow import DataHandle
 from PyConf.tonic import configurable
@@ -18,6 +18,7 @@ from PyConf.application import (default_raw_banks, make_odin,
                                 make_data_with_FetchDataFromFile)
 from PyConf.packing import persistreco_version, default_persistable_location as persistable_location
 from GaudiConf.LbExec import InputProcessTypes, HltSourceID
+import os
 
 
 @configurable
@@ -37,6 +38,41 @@ def tes_root(*, input_process: InputProcessTypes):
         )
 
 
+@configurable
+def upfront_reconstruction(simulation=False):
+    """Return a list DataHandles that define the upfront reconstruction output.
+
+    This differs from `reconstruction` as it should not be used as inputs to
+    other algorithms, but only to define the control flow, i.e. the return
+    value of this function should be ran before all HLT2 lines.
+
+    """
+
+    stream = '/Event/HLT2'
+
+    reco_loc = reco_locations(stream)
+
+    if simulation:
+        reco_loc |= pp2mcp_locations(stream)
+
+    unpackers = [_get_unpacked(v[1], v[0]) for v in reco_loc.values()]
+
+    mc_algs = []
+    if simulation:
+        mc_algs = [
+            get_mc_vertices(os.path.join(stream, 'MC/Vertices')).producer,
+            get_mc_particles(os.path.join(stream, 'MC/Particles')).producer
+        ]
+
+    ### TODO:FIXME take advantage of the fact that the above have datahandles...
+    # i.e. should _not_ have to return decoder here, and should just return the _output handles_ and not the algorithms
+    # i.e. `upfront_reconstruction` should be a drop-in replacement for `reconstruction()`, with the same return type
+    return [
+        dstdata_filter(source='Hlt2'),
+        upfront_decoder(source="Hlt2").producer
+    ] + mc_algs + unpackers
+
+
 @configurable
 def dstdata_filter(source: HltSourceID, raw_banks=default_raw_banks):
     """
@@ -57,7 +93,9 @@ def dstdata_filter(source: HltSourceID, raw_banks=default_raw_banks):
 
 
 @configurable
-def upfront_decoder(*, source: HltSourceID, raw_banks=default_raw_banks):
+def upfront_decoder(*,
+                    source: HltSourceID = "Hlt2",
+                    raw_banks=default_raw_banks):
     """Return a DataHandle for  HltPackedBufferDecoder output
        This is input to all unpacking algorithms except MC and old unpacker
     """
@@ -167,32 +205,39 @@ def get_mc_vertices(location):
 
 
 @configurable
-def reconstruction(*, input_process: InputProcessTypes, obj: str):
+def reconstruction(*,
+                   input_process: InputProcessTypes = "Hlt2",
+                   simulation: bool = False):
     stream = tes_root(input_process=input_process)
     recos = reco_locations(stream=stream)
+    if simulation:
+        recos |= pp2mcp_locations(stream)
+
     reco_locs = {k: _get_unpacked(v[1], v[0]) for k, v in recos.items()}
-    postprocess_unpacked_data(reco_locs)
-    return reco_locs[obj]
+
+    postprocess_unpacked_data(reco_locs, packable=True)
+
+    return reco_locs
 
 
 def get_charged_protoparticles(track_type):
-    return reconstruction(obj=f"{track_type}Protos")
+    return reconstruction()[f"{track_type}Protos"]
 
 
 def get_neutral_protoparticles():
-    return reconstruction(obj="NeutralProtos")
+    return reconstruction()["NeutralProtos"]
 
 
 def get_extended_pvs():
-    return reconstruction(obj="ExtendedPVs")
+    return reconstruction()["ExtendedPVs"]
 
 
 def get_pvs():
-    return reconstruction(obj="PVs")
+    return reconstruction()["PVs"]
 
 
 def get_rec_summary():
-    return reconstruction(obj="RecSummary")
+    return reconstruction()["RecSummary"]
 
 
 def get_mc_track_info():