add an option to specify the max number of files to include in the configuration from the --filein 'das:<dataset name and query>:<number of files' option

Author: vlimant

Configuration/Applications/python/ConfigBuilder.py

@@ -139,7 +139,7 @@ def filesFromList(fileName,s=None):
         print("found parent files:",sec)
     return (prim,sec)
 
-def filesFromDASQuery(query,option="",s=None):
+def filesFromDASQuery(query,option="",s=None,max_files=None):
     import os,time
     import FWCore.ParameterSet.Config as cms
     prim=[]
@@ -172,6 +172,10 @@ def filesFromDASQuery(query,option="",s=None):
     # remove any duplicates
     prim = sorted(list(set(prim)))
     sec = sorted(list(set(sec)))
+    if max_files:
+        max_files=int(max_files)
+        prim = prim[:max_files]
+        sec = sec[:max_files]
     if s:
         if not hasattr(s,"fileNames"):
             s.fileNames=cms.untracked.vstring(prim)
@@ -412,12 +416,19 @@ def addSource(self):
         self.addedObjects.append(("Input source","source"))
 
         def filesFromOption(self):
+            def _datasetname_and_maxfiles(entry):
+                if ":" in entry:
+                    return entry.split(":")
+                else:
+                    return entry,None
+
             for entry in self._options.filein.split(','):
                 print("entry",entry)
                 if entry.startswith("filelist:"):
-                    filesFromList(entry[9:],self.process.source)
+                    filesFromList(entry.split(":",1)[1],self.process.source)
                 elif entry.startswith("dbs:") or entry.startswith("das:"):
-                    filesFromDASQuery('file dataset = %s'%(entry[4:]),self._options.dasoption,self.process.source)
+                    dataset_name,max_files = _datasetname_and_maxfiles(entry.split(":",1)[1])
+                    filesFromDASQuery('file dataset = %s'%(dataset_name),self._options.dasoption,self.process.source,max_files)
                 else:
                     self.process.source.fileNames.append(self._options.dirin+entry)
             if self._options.secondfilein:
@@ -426,9 +437,10 @@ def filesFromOption(self):
                 for entry in self._options.secondfilein.split(','):
                     print("entry",entry)
                     if entry.startswith("filelist:"):
-                        self.process.source.secondaryFileNames.extend((filesFromList(entry[9:]))[0])
+                        self.process.source.secondaryFileNames.extend((filesFromList(entry.split(":",1)[1]))[0])
                     elif entry.startswith("dbs:") or entry.startswith("das:"):
-                        self.process.source.secondaryFileNames.extend((filesFromDASQuery('file dataset = %s'%(entry[4:]),self._options.dasoption))[0])
+                        dataset_name,max_files = _datasetname_and_maxfiles(entry.split(":",1)[1])
+                        self.process.source.secondaryFileNames.extend((filesFromDASQuery('file dataset = %s'%(dataset_name),self._options.dasoption))[0])
                     else:
                         self.process.source.secondaryFileNames.append(self._options.dirin+entry)