Merge pull request #62 from pigay/recursive-dput-dget

Recursive options for dput and dget

Interfaces/Utilities/DCommands.py

@@ -47,7 +47,7 @@ def listFormatPretty( summaries, headers = None, sortKeys = None ):
 
   records = []
   for _k, i in sortKeys:
-    records.append( [ str(x) for x in summaries[i] ] )
+    records.append( [ str( x ) for x in summaries[i] ] )
 
   output = printTable( headers, records, numbering = False, printOut = False, columnSeparator = '  ' )
   return output
@@ -90,7 +90,7 @@ def __init__( self, outputFormat ):
 
   def listFormat( self, list_, headers, sort = None ):
     if self.outputFormat not in self.fmts:
-      return S_ERROR( "ArrayFormatter: Output format not supported: %s not in %s" %
+      return S_ERROR( "ArrayFormatter: Output format not supported: %s not in %s" % 
                       ( self.outputFormat, self.fmts.keys() ) )
 
     if headers is None:
@@ -556,7 +556,7 @@ def guessConfigFromCS( config, section, userName, groupName ):
         host = gConfig.getValue( "/Resources/StorageElements/%s/AccessProtocol.1/Host" % seSite )
         if host and host == voDefaultSEHost:
           # check if SE has rw access
-          retVal = gConfig.getOptionsDict( "/Resources/StorageElements/%s" %
+          retVal = gConfig.getOptionsDict( "/Resources/StorageElements/%s" % 
                                           seSite )
           if retVal[ "OK" ]:
             od = retVal[ "Value" ]
@@ -572,7 +572,7 @@ def guessConfigFromCS( config, section, userName, groupName ):
           break
 
       if defaultSESite:
-        #write to config
+        # write to config
         config.set( section, "default_se", defaultSESite )
 
 def sessionFromProxy( config = DConfig(), sessionDir = None ):
@@ -660,6 +660,9 @@ def getMeta( self, path ):
     if self.isDir( path ):
       return self.catalog.getDirectoryUserMetadata( path )
     return self.catalog.getFileUserMetadata( path )
+
+  def findFilesByMetadata( self, metaDict, path ):
+    return self.catalog.findFilesByMetadata( metaDict, path )
 
 def pathFromArgument( session, arg ):
   path = os.path.normpath( arg )

Interfaces/scripts/dget.py

@@ -17,73 +17,102 @@
 from COMDIRAC.Interfaces import ConfigCache
 from DIRAC.Core.Base import Script
 
+class Params:
+  def __init__ ( self ):
+    self.recursive = False
+
+  def setRecursive( self, arg = None ):
+    self.recursive = True
+
+  def getRecursive( self ):
+    return self.recursive
+
+
+params = Params()
+
+
 Script.setUsageMessage( '\n'.join( [ __doc__.split( '\n' )[1],
                                      'Usage:',
-                                     '  %s lfn... [local_path]' % Script.scriptName,
+                                     '  %s lfn... [local_dir]' % Script.scriptName,
                                      'Arguments:',
                                      ' lfn:          file to download',
-                                     ' local_path:   destination directory',
+                                     ' local_dir:   destination directory',
                                        '', 'Examples:',
                                        '  $ dget ./some_lfn_file /tmp',
                                        ] )
                         )
 
+Script.registerSwitch( "r", "recursive", "recursively get contents of lfn", params.setRecursive )
+
 configCache = ConfigCache()
 Script.parseCommandLine( ignoreErrors = True )
 configCache.cacheConfig()
 
 args = Script.getPositionalArgs()
 
-session = DSession( )
-catalog = DCatalog( )
+session = DSession()
+catalog = DCatalog()
 
 from DIRAC.Interfaces.API.Dirac  import Dirac
 
 dirac = Dirac()
 
 if len( args ) < 1:
   error( "\nError: not enough arguments provided\n%s:" % Script.scriptName )
-  Script.showHelp( )
+  Script.showHelp()
   DIRAC.exit( -1 )
 
 # lfn
-lfn = pathFromArgument( session, args[ 0 ] )
+lfn = pathFromArgument( session, args[0] )
 
-# default local_path: same file name as lfn.
-local_path = os.path.basename( lfn )
-# STRANGE: dirac only accepts directories for download destination
-#pairs = [ ( lfn, local_path ) ]
-pairs = [ ( lfn, os.getcwd( ) ) ]
+localDir = os.getcwd()
+lfns = [( lfn, localDir )]
 
 if len( args ) > 1:
-  # local_path provided must be last argument
-  local_path = args[ -1 ]
-  lfns = args[ :-1 ]
-  pairs = [ ]
+  # localDir provided must be last argument
+  localDir = args[-1]
+  lfns = [( pathFromArgument( session, lfn ), localDir ) for lfn in args[:-1]]
 
-  # STRANGE: dirac only accepts directories for download destination
-  if not os.path.isdir( local_path ):
+  if not os.path.isdir( localDir ):
     critical( "Error: Destination local path must be a directory", -1 )
 
-  if os.path.isdir( local_path ):
-    # we can accept one ore more lfns
-    for lfn in lfns:
-      # STRANGE: dirac only accepts directories for download destination
-      #pairs.append( (pathFromArgument( session, lfn ), os.path.join( local_path, os.path.basename( lfn )) ))
-      pairs.append( (pathFromArgument( session, lfn ), local_path ))
-  else:
-    if len( lfns ) > 1:
-      critical( "Error: Destination path must be a directory when downloading multiple local files", -1 )
-
-    # local filename replace lfn filename
-    pairs.append( (pathFromArgument( session, lfn ), local_path ))
 exitCode = 0
-errmsgs = []
 
-for lfn, local_path in pairs:
-  ret = dirac.getFile( lfn, local_path )
+if params.getRecursive():
+  newLFNs = []
+  for lfn, localDir in lfns:
+    # make sure lfn is an existing directory
+    if not catalog.isDir( lfn ):
+      if catalog.isFile( lfn ):
+        # lfn is a file: simply add it to the list
+        newLFNs.append( ( lfn, localDir ) )
+        continue
+      exitCode = -1
+      error( 'Invalid path: \'%s\'' % lfn )
+      continue
+
+    retVal = catalog.findFilesByMetadata( {}, lfn )
+
+    if not retVal['OK']:
+      exitCode = -2
+      error( retVal['Message'] )
+      continue
+
+    # compute new local destination for subtree files 
+    lfnDirname = os.path.dirname( lfn )
+    for newLFN in retVal['Value']:
+      newLocalDir = os.path.dirname( os.path.join( localDir, os.path.relpath( newLFN, lfnDirname ) ) )
+      newLFNs.append( ( newLFN, newLocalDir ) )
+
+  lfns = newLFNs
+for lfn, localDir in lfns:
+  if params.getRecursive():
+    if not os.path.exists( localDir ):
+      os.makedirs( localDir )
+
+  ret = dirac.getFile( lfn, localDir )
   if not ret['OK']:
-    exitCode = -2
+    exitCode = -3
     error( 'ERROR: %s' % ret['Message'] )
 
 DIRAC.exit( exitCode )

Interfaces/scripts/dput.py

@@ -25,6 +25,7 @@
   class Params:
     def __init__ ( self ):
       self.destinationSE = False
+      self.recursive = False
 
     def setDestinationSE( self, arg ):
       self.destinationSE = arg
@@ -33,9 +34,15 @@ def setDestinationSE( self, arg ):
     def getDestinationSE( self ):
       return self.destinationSE
 
+    def setRecursive( self, arg = None ):
+      self.recursive = True
+
+    def getRecursive( self ):
+      return self.recursive
+
   params = Params()
 
-  Script.setUsageMessage( '\n'.join( [ __doc__.split( '\n' )[1],
+  Script.setUsageMessage( '\n'.join( [__doc__.split( '\n' )[1],
                                        'Usage:',
                                        '  %s [options] local_path[... lfn]' % Script.scriptName,
                                        'Arguments:',
@@ -47,6 +54,7 @@ def getDestinationSE( self ):
                                        ] )
                           )
   Script.registerSwitch( "D:", "destination-se=", "Storage Element where to put replica", params.setDestinationSE )
+  Script.registerSwitch( "r", "recursive", "recursively put contents of local_path", params.setRecursive )
 
 
   configCache = ConfigCache()
@@ -68,46 +76,58 @@ def getDestinationSE( self ):
     DIRAC.exit( 0 )
 
   # local file
-  local_path = args[ 0 ]
+  localPath = args[0]
 
-  # default lfn: same file name as local_path
-  lfn = pathFromArgument( session, os.path.basename( local_path ) )
+  # default lfn: same file name as localPath
+  lfn = pathFromArgument( session, os.path.basename( localPath ) )
 
-  pairs = [ ( local_path, lfn ) ]
+  pairs = [( localPath, lfn )]
 
   if len( args ) > 1:
     # lfn provided must be last argument
-    lfn = pathFromArgument( session, args[ -1 ] )
-    local_paths = args[ :-1 ]
-    pairs = [ ]
+    lfn = pathFromArgument( session, args[-1] )
+    localPaths = args[:-1]
+    pairs = []
 
     if catalog.isDir( lfn ):
       # we can accept one ore more local files
-      for lp in local_paths:
+      for lp in localPaths:
         pairs.append( ( lp, os.path.join( lfn, os.path.basename( lp ) ) ) )
     else:
-      if len( local_paths ) > 1:
+      if len( localPaths ) > 1:
         critical( "Error: Destination LFN must be a directory when registering multiple local files" )
 
       # lfn filename replace local filename
-      pairs.append( ( local_path, lfn ) )
+      pairs.append( ( localPath, lfn ) )
 
   # destination SE
   se = params.getDestinationSE()
   if not se:
     retVal = session.getEnv( "default_se", "DIRAC-USER" )
-    if not retVal[ "OK" ]:
-      error( retVal[ "Message" ] )
-    se = retVal[ "Value" ]
+    if not retVal["OK"]:
+      error( retVal["Message"] )
+    se = retVal["Value"]
 
   exitCode = 0
 
-  for local_path, lfn in pairs:
-    ret = dirac.addFile( lfn, local_path, se, printOutput = False )
-
-  if not ret['OK']:
-    exitCode = -2
-    error( lfn + ': ' + ret['Message'] )
+  if params.getRecursive():
+    newPairs = []
+    for localPath, lfn in pairs:
+      if os.path.isdir( localPath ):
+        for path, _subdirs, files in os.walk( localPath ):
+          newLFNDir = os.path.normpath( os.path.join( lfn, os.path.relpath( path, localPath ) ) )
+          for f in files:
+            pairs.append( ( os.path.join( path, f ), os.path.join( newLFNDir, f ) ) )
+      else:
+        newPairs.append( ( localPath, lfn ) )
+    pairs = newPairs
+
+  for localPath, lfn in pairs:
+    ret = dirac.addFile( lfn, localPath, se, printOutput = False )
+
+    if not ret['OK']:
+      exitCode = -2
+      error( lfn + ': ' + ret['Message'] )
 
 DIRAC.exit( exitCode )