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