Replaced strings of collection protocols with enums

All collection protocols are strings, they are now replaced by Enums

Author: vshekar

config_params.py

@@ -156,3 +156,29 @@ def get_text(cls, enum_value):
 
 OPHYD_COLLECTIONS = {"amx": False, "fmx": False, "nyx": True}
 
+class CollectionProtocols(str, Enum):
+    STANDARD = "standard"
+    RASTER = "raster"
+    VECTOR = "vector"
+    BURN = "burn"
+    RASTER_SCREEN = "rasterScreen"
+    STEP_RASTER = "stepRaster"
+    STEP_VECTOR = "stepVector"
+    MULTI_COL = "multiCol"
+    MULTI_COL_Q = "multiColQ"
+    CHARACTERIZE = "characterize"
+    EDNA_COL = "ednaCol"
+    E_SCAN = "eScan"
+
+    @classmethod
+    def get_beamline_options(cls, beamline):
+        all_protocols = (cls.STANDARD, cls.RASTER, cls.VECTOR, cls.BURN, cls.RASTER_SCREEN,
+                         cls.STEP_RASTER, cls.STEP_VECTOR, cls.MULTI_COL, cls.CHARACTERIZE,
+                         cls.EDNA_COL)
+        beamline_options = {
+            "nyx": (cls.STANDARD, cls.RASTER, cls.VECTOR),
+            "amx": all_protocols,
+            "fmx": all_protocols
+        }
+
+        return beamline_options.get(beamline, ())

daq_lib.py

@@ -605,29 +605,29 @@ def collectData(currentRequest):
     os.system(comm_s)
   logger.debug('starting initial motions - transmission and detector distance')
   daq_macros.setTrans(attenuation)
-  if not prot in ["eScan"]:
+  if not prot in [CollectionProtocols.E_SCAN]:
     beamline_lib.mvaDescriptor("detectorDist",colDist)  
   logger.debug('transmission and detector distance done')
   # now that the detector is in the correct position, get the beam center
   currentRequest['request_obj']['xbeam'] = getPvDesc('beamCenterX')
   currentRequest['request_obj']['ybeam'] = getPvDesc('beamCenterY')
   db_lib.updateRequest(currentRequest)
-  if (prot == "raster"):
+  if prot == CollectionProtocols.RASTER:
     logger.info('entering raster')
     RE(daq_macros.snakeRaster(currentRequest["uid"]))
     status = 0
     logger.info('exiting raster')
-  elif (prot == "stepRaster"):
+  elif prot == CollectionProtocols.STEP_RASTER:
     status = daq_macros.snakeStepRaster(currentRequest["uid"])    
-  elif (prot == "vector" or prot == "stepVector"):
+  elif prot in (CollectionProtocols.VECTOR, CollectionProtocols.STEP_VECTOR):
     imagesAttempted = collect_detector_seq_hw(sweep_start,range_degrees,img_width,exposure_period,file_prefix,data_directory_name,file_number_start,currentRequest)
-  elif (prot == "multiCol"):
+  elif prot == CollectionProtocols.MULTI_COL:
     RE(daq_macros.snakeRaster(currentRequest["uid"]))
-  elif (prot == "rasterScreen"):
+  elif prot == CollectionProtocols.RASTER_SCREEN:
     daq_macros.rasterScreen(currentRequest)    
-  elif (prot == "multiColQ"):
+  elif prot == CollectionProtocols.MULTI_COL_Q:
     daq_macros.multiCol(currentRequest)
-  elif (prot == "eScan"):
+  elif prot == CollectionProtocols.E_SCAN:
     daq_macros.eScan(currentRequest)
   else: #standard, screening, or edna - these may require autoalign, checking first
     if (reqObj["pos_x"] != -999):
@@ -654,7 +654,7 @@ def collectData(currentRequest):
           sweep_start = reqObj["sweep_start"]
       daq_macros.setTrans(attenuation)      
 
-    if (reqObj["protocol"] == "characterize" or reqObj["protocol"] == "ednaCol"):
+    if (reqObj["protocol"] in (CollectionProtocols.CHARACTERIZE, CollectionProtocols.EDNA_COL):
       characterizationParams = reqObj["characterizationParams"]
       index_success = daq_macros.dna_execute_collection3(0.0,img_width,2,exposure_period,data_directory_name+"/",file_prefix,1,-89.0,1,currentRequest)
       if (index_success):        
@@ -693,7 +693,7 @@ def collectData(currentRequest):
           runNum = db_lib.incrementSampleRequestCount(sampleID)
           newReqObj["runNum"] = runNum
           newStratRequest = db_lib.addRequesttoSample(sampleID,newReqObj["protocol"],daq_utils.owner,newReqObj,priority=0,proposalID=daq_utils.getProposalID())
-          if (reqObj["protocol"] == "ednaCol"):
+          if (reqObj["protocol"] == CollectionProtocols.EDNA_COL):
             logger.info("new strat req = ")
             logger.info(newStratRequest)
             db_lib.updatePriority(currentRequest["uid"],-1)
@@ -709,7 +709,7 @@ def collectData(currentRequest):
           beamline_lib.mvaDescriptor("omega",sweep_start)
       collect_detector_seq_hw(sweep_start,range_degrees,img_width,exposure_period,file_prefix,data_directory_name,file_number_start,currentRequest)
   try:
-    if (logMe) and prot == 'raster':
+    if (logMe) and prot == CollectionProtocols.RASTER:
       logMxRequestParams(currentRequest,wait=False)
     elif (logMe):
       logMxRequestParams(currentRequest)
@@ -721,9 +721,9 @@ def collectData(currentRequest):
     logger.error('caught key error in logging: %s' % e)
 
   # collection finished, start processing
-  if reqObj["protocol"] in ("standard", "vector", "raster"):
+  if reqObj["protocol"] in (CollectionProtocols.STANDARD, CollectionProtocols.VECTOR, CollectionProtocols.RASTER):
     send_kafka_message(topic=f'{daq_utils.beamline}.lsdc.documents', event='stop', uuid=currentRequest['uid'], protocol=reqObj["protocol"])
-  if (prot == "vector" or prot == "standard" or prot == "stepVector"):
+  if prot in (CollectionProtocols.VECTOR, CollectionProtocols.STANDARD, CollectionProtocols.STEP_VECTOR):
     if daq_utils.beamline != "nyx":
       seqNum = flyer.detector.cam.sequence_id.get()
       comm_s = os.environ["LSDCHOME"] + "/runSpotFinder4syncW.py " + data_directory_name + " " + file_prefix + " " + str(currentRequest["uid"]) + " " + str(seqNum) + " " + str(currentIspybDCID)+ "&"
@@ -785,7 +785,7 @@ def collect_detector_seq_hw(sweep_start,range_degrees,image_width,exposure_perio
   reqObj = currentRequest["request_obj"]
   protocol = str(reqObj["protocol"])
   sweep_start = sweep_start % 360.0 if sweep_start > 0 else sweep_start % -360
-  if (protocol == "vector" or protocol == "stepVector"):
+  if protocol in (CollectionProtocols.VECTOR, CollectionProtocols.STEP_VECTOR):
     beamline_lib.mvaDescriptor("omega",sweep_start)
   if (image_width == 0):
     number_of_images = range_degrees
@@ -807,22 +807,23 @@ def collect_detector_seq_hw(sweep_start,range_degrees,image_width,exposure_perio
 
   if OPHYD_COLLECTIONS[daq_utils.beamline]:
       logger.info("ophyd collections enabled")
-      if (protocol == "standard"):
+      if (protocol == CollectionProtocols.STANDARD):
         RE(daq_macros.standard_plan_wrapped(currentRequest))
-      elif (protocol == "vector"):
+      elif (protocol == CollectionProtocols.VECTOR):
         RE(daq_macros.vector_plan_wrapped(currentRequest))
   else:  
-    if (protocol == "standard" or protocol == "characterize" or protocol == "ednaCol" or protocol == "burn"):
+    if (protocol in (CollectionProtocols.STANDARD, CollectionProtocols.CHARACTERIZE,
+                     CollectionProtocols.EDNA_COL, CollectionProtocols.BURN):
       logger.info("vectorSync " + str(time.time()))    
       daq_macros.vectorSync()
       logger.info("zebraDaq " + str(time.time()))
 
       vector_params = daq_macros.gatherStandardVectorParams()
       logger.debug(f"vector_params: {vector_params}") 
       RE(daq_macros.standard_zebra_plan(flyer,angleStart,number_of_images,range_degrees,image_width,exposure_period,file_prefix_minus_directory,data_directory_name,file_number, vector_params, file_prefix_minus_directory))
-    elif (protocol == "vector"):
+    elif (protocol == CollectionProtocols.VECTOR):
       RE(daq_macros.vectorZebraScan(currentRequest))
-    elif (protocol == "stepVector"):
+    elif (protocol == CollectionProtocols.STEP_VECTOR):
       daq_macros.vectorZebraStepScan(currentRequest)
     else:
       pass

daq_macros.py

@@ -1431,8 +1431,9 @@ def snakeRasterNoTile(rasterReqID,grain=""):
   rasterRequest["request_obj"]["rasterDef"]["status"] = 2
   protocol = reqObj["protocol"]
   logger.info("protocol = " + protocol)
-  if (protocol == "multiCol" or parentReqProtocol == "multiColQ"):
-    if (parentReqProtocol == "multiColQ"):    
+  if (protocol == CollectionProtocols.MULTI_COL or 
+      parentReqProtocol == CollectionProtocols.MULTI_COL_Q):
+    if (parentReqProtocol == CollectionProtocols.MULTI_COL_Q):    
       multiColThreshold  = parentReqObj["diffCutoff"]
     else:
       multiColThreshold  = reqObj["diffCutoff"]         
@@ -1789,7 +1790,7 @@ def snakeRasterNormal(rasterReqID,grain=""):
                                                             rasterReqID))
         spotFindThread.start()
         spotFindThreadList.append(spotFindThread)
-      send_kafka_message(f'{daq_utils.beamline}.lsdc.documents', event='event', uuid=rasterReqID, protocol="raster", row=i, proc_flag=procFlag)
+      send_kafka_message(f'{daq_utils.beamline}.lsdc.documents', event='event', uuid=rasterReqID, protocol=CollectionProtocols.RASTER, row=i, proc_flag=procFlag)
 
 
     """governor transitions:
@@ -1824,8 +1825,9 @@ def snakeRasterNormal(rasterReqID,grain=""):
       rasterResult = generateGridMap(rasterRequest)
 
       logger.info(f'protocol = {reqObj["protocol"]}')
-      if (reqObj["protocol"] == "multiCol" or parentReqProtocol == "multiColQ"):
-        if (parentReqProtocol == "multiColQ"):    
+      if (reqObj["protocol"] == CollectionProtocols.MULTI_COL or 
+          parentReqProtocol == CollectionProtocols.MULTI_COL_Q):
+        if (parentReqProtocol == CollectionProtocols.MULTI_COL_Q):    
           multiColThreshold  = parentReqObj["diffCutoff"]
         else:
           multiColThreshold  = reqObj["diffCutoff"]         
@@ -2169,7 +2171,7 @@ def snakeRasterBluesky(rasterReqID, grain=""):
                                                               rasterReqID))
           spotFindThread.start()
           spotFindThreadList.append(spotFindThread)
-        send_kafka_message(f'{daq_utils.beamline}.lsdc.documents', event='event', uuid=rasterReqID, protocol="raster", row=row_index, proc_flag=procFlag)
+        send_kafka_message(f'{daq_utils.beamline}.lsdc.documents', event='event', uuid=rasterReqID, protocol=CollectionProtocols.RASTER, row=row_index, proc_flag=procFlag)
         logger.info('row complete')
     """governor transitions:
     initiate transitions here allows for GUI sample/heat map image to update
@@ -2214,8 +2216,9 @@ def snakeRasterBluesky(rasterReqID, grain=""):
       rasterResult = generateGridMap(rasterRequest)
 
       logger.info(f'protocol = {reqObj["protocol"]}')
-      if (reqObj["protocol"] == "multiCol" or parentReqProtocol == "multiColQ"):
-        if (parentReqProtocol == "multiColQ"):
+      if (reqObj["protocol"] == CollectionProtocols.MULTI_COL or 
+          parentReqProtocol == CollectionProtocols.MULTI_COL_Q):
+        if (parentReqProtocol == CollectionProtocols.MULTI_COL_Q):    
           multiColThreshold  = parentReqObj["diffCutoff"]
         else:
           multiColThreshold  = reqObj["diffCutoff"]         
@@ -2820,7 +2823,7 @@ def defineRectRaster(currentRequest,raster_w_s,raster_h_s,stepsizeMicrons_s,xoff
 
   tempnewRasterRequest = daq_utils.createDefaultRequest(sampleID, basePath=currentRequest["request_obj"]["basePath"])
   reqObj = tempnewRasterRequest["request_obj"]
-  reqObj["protocol"] = "raster"
+  reqObj["protocol"] = CollectionProtocols.RASTER
   reqObj["exposure_time"] = getBlConfig("rasterDefaultTime")
   reqObj["img_width"] = getBlConfig("rasterDefaultWidth")
   reqObj["attenuation"] = getBlConfig("rasterDefaultTrans")
@@ -4161,7 +4164,7 @@ def zebraDaqRasterBluesky(flyer, angle_start, num_images, scanWidth, imgWidth, e
                    x_end_um=x_vec_end, y_end_um=y_vec_end, z_end_um=z_vec_end, \
                    file_prefix=filePrefix, data_directory_name=data_directory_name,\
                    detector_dead_time=detectorDeadTime, scan_encoder=scanEncoder, change_state=changeState,\
-                   row_index=row_index, transmission=1, protocol="raster")
+                   row_index=row_index, transmission=1, protocol=CollectionProtocols.RASTER)
     yield from bp.fly([raster_flyer])
 
     logger.info("vector Done")

daq_utils.py

@@ -13,6 +13,7 @@
   logger.error("daq_utils: ISPYB import error, %s" %e)
 
 import db_lib
+from config_params import CollectionProtocols
 
 global beamline
 beamline = os.environ["BEAMLINE_ID"]
@@ -222,7 +223,7 @@ def createDefaultRequest(sample_id,createVisit=True, basePath=None):
                "sweep_start": screenPhist,  "sweep_end": screenPhiend,
                "img_width": screenWidth,
                "exposure_time": screenExptime,
-               "protocol": "standard",
+               "protocol": CollectionProtocols.STANDARD,
                "detDist": screenDist,
                "parentReqID": -1,
                "basePath": basePath,

db_lib.py

@@ -5,6 +5,7 @@
 from collections import defaultdict
 
 import amostra.client.commands as acc
+from config_params import CollectionProtocols
 import conftrak.client.commands as ccc
 import conftrak.exceptions
 import six
@@ -887,6 +888,6 @@ def deleteCompletedRequestsforSample(sid):
   requestList=getRequestsBySampleID(sid)
   for i in range (0,len(requestList)):
     if (requestList[i]["priority"] == -1): #good to clean up completed requests after unmount
-      if (requestList[i]["protocol"] == "raster" or requestList[i]["protocol"] == "vector"):
+      if requestList[i]["protocol"] in (CollectionProtocols.RASTER, CollectionProtocols.VECTOR):
         deleteRequest(requestList[i]['uid'])