sdk/python/arvados/stream.py

   1 from __future__ import print_function
   2 from __future__ import absolute_import
   3 from builtins import object
   4 import collections
   5 import hashlib
   6 import os
   7 import re
   8 import threading
   9 import functools
  10 import copy
  11
  12 from ._ranges import locators_and_ranges, Range
  13 from .arvfile import StreamFileReader
  14 from arvados.retry import retry_method
  15 from .keep import *
  16 from . import config
  17 from . import errors
  18 from ._normalize_stream import normalize_stream
  19
  20 class StreamReader(object):
  21     def __init__(self, tokens, keep=None, debug=False, _empty=False,
  22                  num_retries=0):
  23         self._stream_name = None
  24         self._data_locators = []
  25         self._files = collections.OrderedDict()
  26         self._keep = keep
  27         self.num_retries = num_retries
  28
  29         streamoffset = 0
  30
  31         # parse stream
  32         for tok in tokens:
  33             if debug: print('tok', tok)
  34             if self._stream_name is None:
  35                 self._stream_name = tok.replace('\\040', ' ')
  36                 continue
  37
  38             s = re.match(r'^[0-9a-f]{32}\+(\d+)(\+\S+)*$', tok)
  39             if s:
  40                 blocksize = int(s.group(1))
  41                 self._data_locators.append(Range(tok, streamoffset, blocksize, 0))
  42                 streamoffset += blocksize
  43                 continue
  44
  45             s = re.search(r'^(\d+):(\d+):(\S+)', tok)
  46             if s:
  47                 pos = int(s.group(1))
  48                 size = int(s.group(2))
  49                 name = s.group(3).replace('\\040', ' ')
  50                 if name not in self._files:
  51                     self._files[name] = StreamFileReader(self, [Range(pos, 0, size, 0)], name)
  52                 else:
  53                     filereader = self._files[name]
  54                     filereader.segments.append(Range(pos, filereader.size(), size))
  55                 continue
  56
  57             raise errors.SyntaxError("Invalid manifest format")
  58
  59     def name(self):
  60         return self._stream_name
  61
  62     def files(self):
  63         return self._files
  64
  65     def all_files(self):
  66         return list(self._files.values())
  67
  68     def size(self):
  69         n = self._data_locators[-1]
  70         return n.range_start + n.range_size
  71
  72     def locators_and_ranges(self, range_start, range_size):
  73         return locators_and_ranges(self._data_locators, range_start, range_size)
  74
  75     @retry_method
  76     def _keepget(self, locator, num_retries=None):
  77         return self._keep.get(locator, num_retries=num_retries)
  78
  79     @retry_method
  80     def readfrom(self, start, size, num_retries=None):
  81         """Read up to 'size' bytes from the stream, starting at 'start'"""
  82         if size == 0:
  83             return b''
  84         if self._keep is None:
  85             self._keep = KeepClient(num_retries=self.num_retries)
  86         data = []
  87         for lr in locators_and_ranges(self._data_locators, start, size):
  88             data.append(self._keepget(lr.locator, num_retries=num_retries)[lr.segment_offset:lr.segment_offset+lr.segment_size])
  89         return b''.join(data)
  90
  91     def manifest_text(self, strip=False):
  92         manifest_text = [self.name().replace(' ', '\\040')]
  93         if strip:
  94             for d in self._data_locators:
  95                 m = re.match(r'^[0-9a-f]{32}\+\d+', d.locator)
  96                 manifest_text.append(m.group(0))
  97         else:
  98             manifest_text.extend([d.locator for d in self._data_locators])
  99         manifest_text.extend([' '.join(["{}:{}:{}".format(seg.locator, seg.range_size, f.name.replace(' ', '\\040'))
 100                                         for seg in f.segments])
 101                               for f in list(self._files.values())])
 102         return ' '.join(manifest_text) + '\n'