sdk/python/arvados/commands/put.py

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