21639: Adjust test mocking
[arvados.git] / sdk / python / arvados / _normalize_stream.py
1 # Copyright (C) The Arvados Authors. All rights reserved.
2 #
3 # SPDX-License-Identifier: Apache-2.0
4
5 from __future__ import absolute_import
6 from . import config
7
8 import re
9
10 def escape(path):
11     return re.sub(r'[\\:\000-\040]', lambda m: "\\%03o" % ord(m.group(0)), path)
12
13 def normalize_stream(stream_name, stream):
14     """Take manifest stream and return a list of tokens in normalized format.
15
16     :stream_name:
17       The name of the stream.
18
19     :stream:
20       A dict mapping each filename to a list of `_range.LocatorAndRange` objects.
21
22     """
23
24     stream_name = escape(stream_name)
25     stream_tokens = [stream_name]
26     sortedfiles = list(stream.keys())
27     sortedfiles.sort()
28
29     blocks = {}
30     streamoffset = 0
31     # Go through each file and add each referenced block exactly once.
32     for streamfile in sortedfiles:
33         for segment in stream[streamfile]:
34             if segment.locator not in blocks:
35                 stream_tokens.append(segment.locator)
36                 blocks[segment.locator] = streamoffset
37                 streamoffset += segment.block_size
38
39     # Add the empty block if the stream is otherwise empty.
40     if len(stream_tokens) == 1:
41         stream_tokens.append(config.EMPTY_BLOCK_LOCATOR)
42
43     for streamfile in sortedfiles:
44         # Add in file segments
45         current_span = None
46         fout = escape(streamfile)
47         for segment in stream[streamfile]:
48             # Collapse adjacent segments
49             streamoffset = blocks[segment.locator] + segment.segment_offset
50             if current_span is None:
51                 current_span = [streamoffset, streamoffset + segment.segment_size]
52             else:
53                 if streamoffset == current_span[1]:
54                     current_span[1] += segment.segment_size
55                 else:
56                     stream_tokens.append(u"{0}:{1}:{2}".format(current_span[0], current_span[1] - current_span[0], fout))
57                     current_span = [streamoffset, streamoffset + segment.segment_size]
58
59         if current_span is not None:
60             stream_tokens.append(u"{0}:{1}:{2}".format(current_span[0], current_span[1] - current_span[0], fout))
61
62         if not stream[streamfile]:
63             stream_tokens.append(u"0:0:{0}".format(fout))
64
65     return stream_tokens