17464: Reset database after upload tests
[arvados.git] / sdk / python / arvados / commands / ls.py
1 # Copyright (C) The Arvados Authors. All rights reserved.
2 #
3 # SPDX-License-Identifier: Apache-2.0
4
5 from __future__ import print_function
6 from __future__ import division
7
8 import argparse
9 import collections
10 import logging
11 import re
12 import sys
13
14 import arvados
15 import arvados.commands._util as arv_cmd
16
17 from arvados._version import __version__
18
19 FileInfo = collections.namedtuple('FileInfo', ['stream_name', 'name', 'size'])
20
21 def parse_args(args):
22     parser = argparse.ArgumentParser(
23         description='List contents of a manifest',
24         parents=[arv_cmd.retry_opt])
25
26     parser.add_argument('locator', type=str,
27                         help="""Collection UUID or locator, optionally with a subdir path.""")
28     parser.add_argument('-s', action='store_true',
29                         help="""List file sizes, in KiB.""")
30     parser.add_argument('--version', action='version',
31                         version="%s %s" % (sys.argv[0], __version__),
32                         help='Print version and exit.')
33
34     return parser.parse_args(args)
35
36 def size_formatter(coll_file):
37     return "{:>10}".format((coll_file.size + 1023) // 1024)
38
39 def name_formatter(coll_file):
40     return "{}/{}".format(coll_file.stream_name, coll_file.name)
41
42 def main(args, stdout, stderr, api_client=None, logger=None):
43     args = parse_args(args)
44
45     if api_client is None:
46         api_client = arvados.api('v1')
47
48     if logger is None:
49         logger = logging.getLogger('arvados.arv-ls')
50
51     try:
52         r = re.search(r'^(.*?)(/.*)?$', args.locator)
53         collection = r.group(1)
54         get_prefix = r.group(2)
55
56         cr = arvados.CollectionReader(collection, api_client=api_client,
57                                       num_retries=args.retries)
58         if get_prefix:
59             if get_prefix[-1] == '/':
60                 get_prefix = get_prefix[:-1]
61             stream_name = '.' + get_prefix
62             reader = cr.find(stream_name)
63             if not (isinstance(reader, arvados.CollectionReader) or
64                     isinstance(reader, arvados.collection.Subcollection)):
65                 logger.error("'{}' is not a subdirectory".format(get_prefix))
66                 return 1
67         else:
68             stream_name = '.'
69             reader = cr
70     except (arvados.errors.ApiError,
71             arvados.errors.ArgumentError,
72             arvados.errors.NotFoundError) as error:
73         logger.error("error fetching collection: {}".format(error))
74         return 1
75
76     formatters = []
77     if args.s:
78         formatters.append(size_formatter)
79     formatters.append(name_formatter)
80
81     for f in files_in_collection(reader, stream_name):
82         print(*(info_func(f) for info_func in formatters), file=stdout)
83
84     return 0
85
86 def files_in_collection(c, stream_name='.'):
87     # Sort first by file type, then alphabetically by file path.
88     for i in sorted(c.keys(),
89                     key=lambda k: (
90                         isinstance(c[k], arvados.collection.Subcollection),
91                         k.upper())):
92         if isinstance(c[i], arvados.arvfile.ArvadosFile):
93             yield FileInfo(stream_name=stream_name,
94                            name=i,
95                            size=c[i].size())
96         elif isinstance(c[i], arvados.collection.Subcollection):
97             for f in files_in_collection(c[i], "{}/{}".format(stream_name, i)):
98                 yield f