python/TkAlAllInOneTool/findAndChange.py

0001 import os
0002 from FWCore.ParameterSet.pfnInPath import pfnInPath
0003 import ROOT
0004
0005 ##############################################
0006 def digest_path(path):
0007 ##############################################
0008     """ Ensure that everything is done for path to exist
0009     Arguments:
0010     - path: String that can be both directory and file
0011     Return:
0012     - general: environmental variables are expanded
0013     - directory: is checked to exist
0014     - file: is checked to exist with backup directory being searched in cms-data
0015     """
0016     # sanity check for string argument
0017     if not isinstance(path, str):
0018         return path
0019
0020     path_expanded = os.path.expandvars(path)
0021
0022     # split path in folders
0023     protocol = ""
0024     if "://" in path_expanded:
0025         protocol = path_expanded.split("://")[0]+"://"
0026         path_d = path_expanded.split("://")[1]
0027     elif ":" in path_expanded:
0028         protocol = path_expanded.split(":")[0]+':'
0029         path_d = ":".join(path_expanded.split(":")[1:])
0030         # Similar to just `split(':')[1]`, but handles the case in which the rest of the path contains one or more ':'
0031     else:
0032         path_d = path_expanded
0033
0034     path_s = path_d.split(os.sep)
0035
0036     placeholderIdx = []
0037     for ipart,part in enumerate(path_s):
0038         # Look for {} placeholder to be replaced internally
0039         if "{}" in part:
0040             placeholderIdx.append(ipart)
0041
0042     # re-join folders into full path
0043     # only check path up to first placeholder occurence
0044     if len(placeholderIdx) > 0:
0045         path_to_check = os.path.join(*path_s[:placeholderIdx[0]])
0046         # re add front / if needed
0047         if path_d.startswith(os.sep):
0048             path_to_check = os.sep + path_to_check
0049     else:
0050         path_to_check = path_d
0051
0052     # check for path to exist
0053     if not os.path.exists(path_to_check) and "." in os.path.splitext(path_to_check)[-1]:
0054         # in case of directory pointing to file try backup
0055         _file = pfnInPath(path_to_check)
0056         if "file:" in _file:
0057             return _file.split(":")[-1]
0058
0059     # re add protocol declaration
0060     if protocol != "": path_d = protocol + path_d
0061
0062     # if all is OK return path to directory or file
0063     return path_d
0064
0065 #########################################
0066 def get_root_color(value):
0067 #########################################
0068     """
0069        Returns an integer correspondig to the ROOT color
0070     """
0071     if(isinstance(value, str)):
0072         if(value.isdigit()):
0073             return int(value)
0074         elif('-' in value):
0075             pre, op, post = value.partition('-')
0076             return get_root_color(pre.strip()) - get_root_color(post.strip())
0077         elif('+' in value):
0078             pre, op, post = value.partition('+')
0079             return get_root_color(pre.strip()) + get_root_color(post.strip())
0080         else:
0081             return getattr(ROOT.EColor, value)
0082     else:
0083         return int(value)
0084
0085 #########################################
0086 def get_all_keys(var):
0087 #########################################
0088     """
0089        Generate all keys for nested dictionary
0090        - reserved keywords are not picked up
0091     """
0092     reserved_keys = ["customrighttitle","title","Rlabel"]
0093     if hasattr(var,'items'):
0094         for k, v in var.items():
0095             if k in reserved_keys: continue
0096             if isinstance(v, dict):
0097                 for result in get_all_keys(v):
0098                     yield result
0099             elif isinstance(v, list):
0100                 for d in v:
0101                     for result in get_all_keys(d):
0102                         yield result
0103             else:
0104                 yield k
0105
0106 ####################################################
0107 def find_and_change(keys, var, alt=digest_path):
0108 ####################################################
0109     """Perform effective search for keys in nested dictionary
0110        - if key is found, corresponding value is "digested"
0111        - generator is returned for printout purpose only
0112        - original var is overwritten
0113     """
0114     if hasattr(var,'items'):
0115         if len(keys) == 0:
0116             for key in get_all_keys(var):
0117                 keys.append(key)
0118         for key in keys:
0119             for k, v in var.items():
0120                 if k == key:
0121                     if "color" in key:
0122                         if isinstance(v,list):
0123                             var[k] = [get_root_color(_v) for _v in v]
0124                         else:
0125                             var[k] = get_root_color(v)
0126                     else:
0127                         if isinstance(v,list):
0128                             var[k] = [alt(_v) for _v in v]
0129                         else:
0130                             var[k] = alt(v)
0131                     yield alt(v)
0132                 if isinstance(v, dict):
0133                     for result in find_and_change([key], v):
0134                         yield result
0135                 elif isinstance(v, list):
0136                     for d in v:
0137                         for result in find_and_change([key], d):
0138                             yield result