python/TkAlAllInOneTool/findAndChange.py

0001 import os
0002 from FWCore.ParameterSet.pfnInPath import pfnInPath
0003
0004 ##############################################
0005 def digest_path(path):
0006 ##############################################
0007     """ Ensure that everything is done for path to exist
0008     Arguments:
0009     - path: String that can be both directory and file
0010     Return:
0011     - general: environmental variables are expanded
0012     - directory: is checked to exist
0013     - file: is checked to exist with backup directory being searched in cms-data
0014     """
0015     # sanity check for string argument
0016     if not isinstance(path, str):
0017         return path
0018
0019     # split path in folders
0020     protocol = ""
0021     if "://" in path:
0022         protocol = path.split("://")[0]+"://"
0023         path_s = path.split("://")[1].split(os.sep)
0024     else:
0025         path_s = path.split(os.sep)
0026
0027     path_d_s = []
0028     placeholderIdx = []
0029     for ipart,part in enumerate(path_s):
0030         # Look for environment variables such as $CMSSW_BASE
0031         if part.startswith('$'):
0032             env_var = part[1:].replace('{', '').replace('}', '')
0033             path_d_s.append(os.environ[env_var])
0034         # Look for {} placeholder to be replaced internally
0035         elif "{}" in part:
0036             placeholderIdx.append(ipart)
0037             path_d_s.append(part)
0038         else:
0039             path_d_s.append(part)
0040
0041     # re-join folders into full path
0042     # only check path up to first placeholder occurence
0043     path_d = os.path.join(*path_d_s)
0044     if len(placeholderIdx) > 0:
0045         path_to_check = os.path.join(*path_d_s[:placeholderIdx[0]])
0046     else:
0047         path_to_check = path_d
0048
0049     # re add front / if needed
0050     if path.startswith(os.sep):
0051         path_d = os.sep + path_d
0052         path_to_check = os.sep + path_to_check
0053
0054     # check for path to exist
0055     if not os.path.exists(path_to_check) and "." in os.path.splitext(path_to_check)[-1]:
0056         # in case of directory pointing to file try backup
0057         _file = pfnInPath(path_to_check)
0058         if "file:" in _file:
0059             return _file.split(":")[-1]
0060
0061     # re add protocol declaration
0062     if protocol != "": path_d = protocol + path_d
0063
0064     # if all is OK return path to directory or file
0065     return path_d
0066
0067 #########################################
0068 def get_all_keys(var):
0069 #########################################
0070     """
0071        Generate all keys for nested dictionary
0072     """
0073     if hasattr(var,'items'):
0074         for k, v in var.items():
0075             if isinstance(v, dict):
0076                 for result in get_all_keys(v):
0077                     yield result
0078             elif isinstance(v, list):
0079                 for d in v:
0080                     for result in get_all_keys(d):
0081                         yield result
0082             else:
0083                 yield k
0084
0085 ####################################################
0086 def find_and_change(keys, var, alt=digest_path):
0087 ####################################################
0088     """Perform effective search for keys in nested dictionary
0089        - if key is found, corresponding value is "digested"
0090        - generator is returned for printout purpose only
0091        - original var is overwritten
0092     """
0093     if hasattr(var,'items'):
0094         if len(keys) == 0:
0095             for key in get_all_keys(var):
0096                 keys.append(key)
0097         for key in keys:
0098             for k, v in var.items():
0099                 if k == key:
0100                     if isinstance(v,list):
0101                         var[k] = [alt(_v) for _v in v]
0102                     else:
0103                         var[k] = alt(v)
0104                     yield alt(v)
0105                 if isinstance(v, dict):
0106                     for result in find_and_change([key], v):
0107                         yield result
0108                 elif isinstance(v, list):
0109                     for d in v:
0110                         for result in find_and_change([key], d):
0111                             yield result