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