sdk/python/bin/arv-put

   1 #!/usr/bin/env python
   2
   3 # TODO:
   4 # --md5sum - display md5 of each file as read from disk
   5
   6 import argparse
   7 import os
   8 import sys
   9
  10 parser = argparse.ArgumentParser(
  11     description='Copy data from the local filesystem to Keep.')
  12 parser.add_argument('paths', metavar='path', type=str, nargs='*',
  13                     help="""
  14 Local file or directory. Default: read from standard input.
  15 """)
  16 parser.add_argument('--max-manifest-depth', type=int, metavar='N', default=-1,
  17                     help="""
  18 Maximum depth of directory tree to represent in the manifest
  19 structure. A directory structure deeper than this will be represented
  20 as a single stream in the manifest. If N=0, the manifest will contain
  21 a single stream. Default: -1 (unlimited), i.e., exactly one manifest
  22 stream per filesystem directory that contains files.
  23 """)
  24 group = parser.add_mutually_exclusive_group()
  25 group.add_argument('--as-stream', action='store_true', dest='stream',
  26                    help="""
  27 Synonym for --stream.
  28 """)
  29 group.add_argument('--stream', action='store_true',
  30                    help="""
  31 Store the file content and display the resulting manifest on
  32 stdout. Do not write the manifest to Keep or save a Collection object
  33 in Arvados.
  34 """)
  35 group.add_argument('--as-manifest', action='store_true', dest='manifest',
  36                    help="""
  37 Synonym for --manifest.
  38 """)
  39 group.add_argument('--in-manifest', action='store_true', dest='manifest',
  40                    help="""
  41 Synonym for --manifest.
  42 """)
  43 group.add_argument('--manifest', action='store_true',
  44                    help="""
  45 Store the file data and resulting manifest in Keep, save a Collection
  46 object in Arvados, and display the manifest locator (Collection uuid)
  47 on stdout. This is the default behavior.
  48 """)
  49 group.add_argument('--as-raw', action='store_true', dest='raw',
  50                    help="""
  51 Synonym for --raw.
  52 """)
  53 group.add_argument('--raw', action='store_true',
  54                    help="""
  55 Store the file content and display the data block locators on stdout,
  56 separated by commas, with a trailing newline. Do not store a
  57 manifest.
  58 """)
  59 parser.add_argument('--use-filename', type=str, default=None, dest='filename',
  60                     help="""
  61 Synonym for --filename.
  62 """)
  63 parser.add_argument('--filename', type=str, default=None,
  64                     help="""
  65 Use the given filename in the manifest, instead of the name of the
  66 local file. This is useful when "-" or "/dev/stdin" is given as an
  67 input file. It can be used only if there is exactly one path given and
  68 it is not a directory. Implies --manifest.
  69 """)
  70 group = parser.add_mutually_exclusive_group()
  71 group.add_argument('--progress', action='store_true',
  72                    help="""
  73 Display human-readable progress on stderr (bytes and, if possible,
  74 percentage of total data size). This is the default behavior when
  75 stderr is a tty.
  76 """)
  77 group.add_argument('--no-progress', action='store_true',
  78                    help="""
  79 Do not display human-readable progress on stderr, even if stderr is a
  80 tty.
  81 """)
  82 group.add_argument('--batch-progress', action='store_true',
  83                    help="""
  84 Display machine-readable progress on stderr (bytes and, if known,
  85 total data size).
  86 """)
  87
  88 args = parser.parse_args()
  89
  90 if len(args.paths) == 0:
  91     args.paths += ['/dev/stdin']
  92
  93 if len(args.paths) != 1 or os.path.isdir(args.paths[0]):
  94     if args.filename:
  95         parser.error("""
  96 --filename argument cannot be used when storing a directory or
  97 multiple files.
  98 """)
  99
 100 # Turn on --progress by default if stderr is a tty.
 101 if (not (args.batch_progress or args.no_progress)
 102     and os.isatty(sys.stderr.fileno())):
 103     args.progress = True
 104
 105
 106 import arvados
 107 import re
 108 import string
 109
 110 class CollectionWriterWithProgress(arvados.CollectionWriter):
 111     def flush_data(self, *args, **kwargs):
 112         if not getattr(self, 'display_type', None):
 113             return
 114         if not hasattr(self, 'bytes_flushed'):
 115             self.bytes_flushed = 0
 116         self.bytes_flushed += self._data_buffer_len
 117         super(CollectionWriterWithProgress, self).flush_data(*args, **kwargs)
 118         self.bytes_flushed -= self._data_buffer_len
 119         if self.display_type == 'machine':
 120             sys.stderr.write('%s %d: %d written %d total\n' %
 121                              (sys.argv[0],
 122                               os.getpid(),
 123                               self.bytes_flushed,
 124                               getattr(self, 'bytes_expected', -1)))
 125         elif getattr(self, 'bytes_expected', 0) > 0:
 126             pct = 100.0 * self.bytes_flushed / self.bytes_expected
 127             sys.stderr.write('\r%dM / %dM %.1f%% ' %
 128                              (self.bytes_flushed >> 20,
 129                               self.bytes_expected >> 20, pct))
 130         else:
 131             sys.stderr.write('\r%d ' % self.bytes_flushed)
 132     def manifest_text(self, *args, **kwargs):
 133         manifest_text = (super(CollectionWriterWithProgress, self)
 134                          .manifest_text(*args, **kwargs))
 135         if getattr(self, 'display_type', None):
 136             if self.display_type == 'human':
 137                 sys.stderr.write('\n')
 138             self.display_type = None
 139         return manifest_text
 140
 141 if args.progress:
 142     writer = CollectionWriterWithProgress()
 143     writer.display_type = 'human'
 144 elif args.batch_progress:
 145     writer = CollectionWriterWithProgress()
 146     writer.display_type = 'machine'
 147 else:
 148     writer = arvados.CollectionWriter()
 149
 150 if args.paths == ['-']:
 151     args.paths = ['/dev/stdin']
 152     if not args.filename:
 153         args.filename = '-'
 154
 155 # Walk the given directory trees and stat files, adding up file sizes,
 156 # so we can display progress as percent
 157 writer.bytes_expected = 0
 158 for path in args.paths:
 159     if os.path.isdir(path):
 160         for filename in arvados.util.listdir_recursive(path):
 161             writer.bytes_expected += os.path.getsize(
 162                 os.path.join(path, filename))
 163     elif not os.path.isfile(path):
 164         del writer.bytes_expected
 165         break
 166     else:
 167         writer.bytes_expected += os.path.getsize(path)
 168
 169 # Copy file data to Keep.
 170 for path in args.paths:
 171     if os.path.isdir(path):
 172         writer.write_directory_tree(path,
 173                                     max_manifest_depth=args.max_manifest_depth)
 174     else:
 175         writer.start_new_stream()
 176         writer.start_new_file(args.filename or os.path.split(path)[1])
 177         with open(path, 'rb') as f:
 178             while True:
 179                 buf = f.read(2**26)
 180                 if len(buf) == 0:
 181                     break
 182                 writer.write(buf)
 183
 184 if args.stream:
 185     print writer.manifest_text(),
 186 elif args.raw:
 187     writer.finish_current_stream()
 188     print string.join(writer.data_locators(), ',')
 189 else:
 190     # Register the resulting collection in Arvados.
 191     arvados.api().collections().create(
 192         body={
 193             'uuid': writer.finish(),
 194             'manifest_text': writer.manifest_text(),
 195             },
 196         ).execute()
 197
 198     # Print the locator (uuid) of the new collection.
 199     print writer.finish()