Modify the functions in FileFromDirac

To fulfill the requirements in #9, modify FileFromDirac

• Filter certain sites
• Sort the sites for MC

python/PRConfig/FilesFromDirac.py

@@ -14,7 +14,11 @@ import subprocess
 from collections import defaultdict
 
 
-def get_access_urls_mc(bkkpath, evttype, filetypes, max_files=500):
+def get_access_urls_mc(bkkpath,
+                       evttype,
+                       filetypes,
+                       sites_to_remove=[],
+                       max_files=500):
     customEnv = {}
 
     # set custom grid proxy path if exists
@@ -66,17 +70,22 @@ def get_access_urls_mc(bkkpath, evttype, filetypes, max_files=500):
 
     # Get the first URL (if more than one) for each LFN, while skipping
     # LFNs for which we couldn't find an URL (e.g. a site was down).
-    lfns = [urls[lfn][0] for lfn in file_list if lfn in urls]
+    lfns_tmp = [urls[lfn][0] for lfn in file_list if lfn in urls]
 
     # Filter out some failing grid sites/files from the list
-    excluded = ['stfc.ac.uk']
-    return [lfn for site in excluded for lfn in lfns if site not in lfn]
+    excluded = ['stfc.ac.uk'] + sites_to_remove
+    lfns = [
+        lfn for lfn in lfns_tmp if not any(site in lfn for site in excluded)
+    ]
+    lfns = sorted(lfns, key=lambda str: not "eoslhcb.cern.ch" in str)
+
+    return lfns
 
     # TODO warn if some of the first N files was not resolved to a URL
     # since then one would get numerically different results.
 
 
-def get_access_urls_data(bkkpath, max_files=500):
+def get_access_urls_data(bkkpath, sites_to_remove=[], max_files=500):
     customEnv = {}
 
     # set custom grid proxy path if exists
@@ -102,9 +111,6 @@ def get_access_urls_data(bkkpath, max_files=500):
 
     # ensure files are always in same order
     file_list.sort()
-    if any(file.endswith("raw") for file in file_list):
-        raise NotImplementedError(
-            "File ending with 'raw' found, please write a script manually")
 
     print(
         "#### Checking output of . /cvmfs/lhcb.cern.ch/lib/LbEnv --quiet; lb-dirac dirac-dms-lfn-accessURL --Terminal"
@@ -134,12 +140,16 @@ def get_access_urls_data(bkkpath, max_files=500):
 
     # Get the first URL (if more than one) for each LFN, while skipping
     # LFNs for which we couldn't find an URL (e.g. a site was down).
-    lfns = [urls[lfn][0] for lfn in file_list if lfn in urls]
+    lfns_tmp = [urls[lfn][0] for lfn in file_list if lfn in urls]
 
     # Filter out some failing grid sites/files from the list
-    excluded = ['stfc.ac.uk']
-    lfns = [lfn for site in excluded for lfn in lfns if site not in lfn]
+    excluded = ['stfc.ac.uk'] + sites_to_remove
+    lfns = [
+        lfn for lfn in lfns_tmp if not any(site in lfn for site in excluded)
+    ]
     lfns = sorted(lfns, key=lambda str: not "eoslhcb.cern.ch" in str)
+    if any(file.endswith("raw") for file in file_list):
+        lfns = [lfn for lfn in lfns if "eoslhcb.cern.ch" in lfn]
 
     return lfns