projects / arvados.git / blob
summary | shortlog | log | commit | commitdiff | tree
history | raw | HEAD
Merge branch '2756-eventbus-in-workbench' into 2678-owned_by
[arvados.git] / 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()