python-misc-utils 0.2__py3-none-any.whl

py_misc_utils/split.py

@@ -0,0 +1,147 @@
+import collections
+import re
+
+from . import assert_checks as tas
+
+
+class _Skipper:
+
+  def __init__(self, quote_rx):
+    self.quote_rx = quote_rx
+    self.next_pos = 0
+
+  def skip(self, data, pos):
+    next_pos = self.next_pos - pos
+    if next_pos <= 0:
+      m = re.search(self.quote_rx, data)
+      next_pos = m.start() if m else len(data)
+      self.next_pos = pos + next_pos
+
+    return next_pos
+
+
+def _chars_regex(chars):
+  rexs = bytearray(b'[\\')
+  for c in sorted(chars):
+    rexs.extend((ord('\\'), c))
+
+  rexs.append(ord(']'))
+
+  return re.compile(bytes(rexs))
+
+
+def _specials_regex(qmap):
+  return _chars_regex(set(tuple(qmap.keys()) + tuple(qmap.values())))
+
+
+def _split_forward(data, pos, split_rx, skipper, seq):
+  pdata = data[pos:]
+
+  xm = re.search(split_rx, pdata)
+  if xm:
+    seq_pos, next_pos = xm.start(), xm.end()
+  else:
+    seq_pos = next_pos = len(pdata)
+
+  skip_pos = skipper.skip(pdata, pos)
+  if skip_pos < seq_pos:
+    seq_pos = next_pos = skip_pos
+    xm = None
+
+  seq.extend(pdata[: seq_pos])
+
+  return pos + next_pos, xm is not None
+
+
+SplitContext = collections.namedtuple('SplitContext', 'map, quote_rx, quote_sprx')
+
+def make_context(quote_map):
+  bmap = {ord(k): ord(v) for k, v in quote_map.items()}
+
+  return SplitContext(map=bmap,
+                      quote_rx=_chars_regex(bmap.keys()),
+                      quote_sprx=_specials_regex(bmap))
+
+
+def _to_bytes(data, split_rx):
+  if isinstance(data, str):
+    data = data.encode()
+  if isinstance(split_rx, str):
+    split_rx = split_rx.encode()
+
+  split_rx = re.compile(split_rx) if isinstance(split_rx, bytes) else split_rx
+
+  return memoryview(data), split_rx
+
+
+_QUOTE_MAP = {
+  '"': '"',
+  "'": "'",
+  '`': '`',
+  '(': ')',
+  '{': '}',
+  '[': ']',
+  '<': '>',
+}
+_QUOTE_CTX = make_context(_QUOTE_MAP)
+
+_Quote = collections.namedtuple('Quote', 'closec, pos, nest_ok')
+
+def split(data, split_rx, quote_ctx=None):
+  qctx = quote_ctx or _QUOTE_CTX
+
+  bdata, bsplit_rx = _to_bytes(data, split_rx)
+  skipper = _Skipper(qctx.quote_rx)
+
+  sval = ord('\\')
+  pos, qstack, parts, seq = 0, [], [], bytearray()
+  while pos < len(bdata):
+    if seq and seq[-1] == sval:
+      seq.append(bdata[pos])
+      pos += 1
+    elif qstack:
+      m = re.search(qctx.quote_sprx, bdata[pos:])
+      if not m:
+        break
+
+      seq.extend(bdata[pos: pos + m.start()])
+      pos += m.start()
+      c = bdata[pos]
+      tq = qstack[-1]
+      if c == tq.closec:
+        qstack.pop()
+      elif tq.nest_ok and (cc := qctx.map.get(c)) is not None:
+        qstack.append(_Quote(cc, pos, c != cc))
+      seq.append(c)
+      pos += 1
+    else:
+      kpos, is_split = _split_forward(bdata, pos, bsplit_rx, skipper, seq)
+      if is_split:
+        parts.append(seq)
+        seq = bytearray()
+      elif kpos < len(bdata):
+        c = bdata[kpos]
+        if (cc := qctx.map.get(c)) is not None:
+          qstack.append(_Quote(cc, kpos, c != cc))
+        seq.append(c)
+        kpos += 1
+      pos = max(kpos, pos + 1)
+
+  tas.check_eq(len(qstack), 0, msg=f'Unmatched quotes during split: "{data}"\n  {qstack}')
+  if seq or parts:
+    parts.append(seq)
+
+  return tuple(p.decode() for p in parts) if isinstance(data, str) else tuple(parts)
+
+
+def unquote(data, quote_map=None):
+  if len(data) >= 2:
+    quote_map = quote_map or _QUOTE_MAP
+    cc = quote_map.get(data[0])
+    if cc == data[-1]:
+      udata = data[1: -1]
+
+      return re.sub(rf'\\{cc}', rf'{cc}', udata) if cc == data[0] else udata
+
+  return data
+

py_misc_utils/state.py

@@ -0,0 +1,53 @@
+import pickle
+
+from . import gfs
+
+
+_STATE_KEY = '__SB_STATE__'
+
+
+def _kname(cls, name):
+  return f'{cls.__name__}.{name}'
+
+
+class StateBase:
+
+  def _get_state(self, state):
+    return state
+
+  def _set_state(self, state):
+    self.__dict__.update(state)
+
+  def _store_state(self, cls, **kwargs):
+    sdict = getattr(self, _STATE_KEY, None)
+    if sdict is None:
+      sdict = dict()
+      setattr(self, _STATE_KEY, sdict)
+
+    for k, v in kwargs.items():
+      sdict[_kname(cls, k)] = v
+
+  def _load_state(self, cls, state, name, defval=None):
+    sdict = state.get(_STATE_KEY)
+
+    return sdict.get(_kname(cls, name), defval) if sdict is not None else defval
+
+
+def to_state(obj, path):
+  # Needs a copy here, as the _get_state() call chains will modify the state.
+  state = obj._get_state(obj.__dict__.copy())
+  with gfs.open(path, mode='wb') as sfd:
+    pickle.dump(state, sfd)
+
+
+def from_state(cls, path, **kwargs):
+  with gfs.open(path, mode='rb') as sfd:
+    state = pickle.load(sfd)
+
+  state.update(kwargs)
+
+  obj = cls.__new__(cls)
+  obj._set_state(state)
+
+  return obj
+

py_misc_utils/std_module.py

@@ -0,0 +1,56 @@
+import functools
+import importlib
+import os
+import sys
+
+from . import core_utils as cu
+
+
+def _module_origin(modname):
+  module = sys.modules.get(modname)
+  if module is None:
+    try:
+      module = importlib.import_module(modname)
+    except ModuleNotFoundError:
+      pass
+
+  if module is not None:
+    path = getattr(module, '__file__', None)
+    if path is None:
+      spec = getattr(module, '__spec__', None)
+      path = spec.origin if spec is not None else None
+
+    return path
+
+
+def _module_libpath(modname):
+  origin = _module_origin(modname)
+  if origin not in {None, 'built-in'}:
+    lib_path = os.path.dirname(origin)
+
+    return lib_path if lib_path else None
+
+
+# Some of the standard modules. Should be enough to get coverage of the
+# Python standard library path (there are more than one since some might
+# turn "built-in" and not have a __file__ or __spec__).
+_STDLIB_MODULES = (
+  'abc',
+  'copy',
+  'io',
+  'os',
+  'pickle',
+  'random',
+  'string',
+  'types',
+)
+_STDLIB_PATHS = set(filter(lambda x: x is not None,
+                           (_module_libpath(m) for m in _STDLIB_MODULES)))
+
+@functools.cache
+def is_std_module(modname):
+  modname = cu.root_module(modname)
+  lib_path = _module_libpath(modname)
+
+  return lib_path is None or lib_path in _STDLIB_PATHS
+

py_misc_utils/stream_dataframe.py

@@ -0,0 +1,176 @@
+import bisect
+import collections
+
+import numpy as np
+import pandas as pd
+
+from . import assert_checks as tas
+from . import np_utils as npu
+from . import tensor_stream as ts
+from . import utils as ut
+
+
+WriteField = collections.namedtuple('WriteField', 'dtype')
+
+
+class StreamDataWriter:
+
+  def __init__(self, fields, path):
+    self._writer = ts.Writer(path)
+    self._fields = collections.OrderedDict()
+    if isinstance(fields, str):
+      sfields = tuple(tuple(ut.resplit(x, '=')) for x in ut.comma_split(fields))
+    else:
+      sfields = fields
+
+    for field, dtype in sfields:
+      self._fields[field] = WriteField(dtype=np.dtype(dtype))
+
+  # Note that the tensors handed over to the write() API will become owned by
+  # the StreamDataWriter obect, and cannot be written over after the write operation.
+  def write(self, **kwargs):
+    args = []
+    for field, wfield in self._fields.items():
+      data = kwargs.get(field)
+      tas.check_is_not_none(data, msg=f'Missing "{field}" data in write operation')
+
+      if isinstance(data, np.ndarray):
+        if data.dtype != wfield.dtype:
+          data = data.astype(wfield.dtype)
+      else:
+        data = np.array(data, dtype=wfield.dtype)
+
+      args.append(data)
+
+    self._writer.write(*args)
+
+  def write_dataframe(self, df):
+    wargs = collections.OrderedDict()
+    for field in self._fields.keys():
+      wargs[field] = df[field].to_numpy()
+
+    self.write(**wargs)
+
+  def flush(self):
+    state = dict(fields=self._fields)
+
+    self._writer.flush(state=state)
+
+
+class StreamDataReader:
+
+  def __init__(self, path):
+    self._reader = ts.Reader(path)
+    self._fields = self._reader.state['fields']
+    self._fields_id = {field: i for i, field in enumerate(self._fields.keys())}
+
+  def __len__(self):
+    return len(self._reader)
+
+  def fields(self):
+    return tuple(self._fields.keys())
+
+  @property
+  def dtype(self):
+    return tuple(wfield.dtype for wfield in self._fields.values())
+
+  def get_slice(self, start, size=None):
+    data = collections.OrderedDict()
+    for i, field in enumerate(self._fields.keys()):
+      data[field] = self._reader.get_slice(i, start, size=size)
+
+    return data
+
+  def get_field_slice(self, field, start, size=None):
+    fid = self._fields_id[field]
+
+    return self._reader.get_slice(fid, start, size=size)
+
+  def typed_fields(self):
+    return tuple((field, wfield.dtype) for field, wfield in self._fields.items())
+
+  def empty_array(self, size):
+    rdata = collections.OrderedDict()
+    for field, dtype in self.typed_fields():
+      if npu.is_numeric(dtype):
+        rdata[field] = np.empty(size, dtype=dtype)
+      else:
+        rdata[field] = [None] * size
+
+    return rdata
+
+
+def _compute_indices(reader, field, start=None, end=None, reverse=False):
+  fvalues = reader.get_field_slice(field, 0)
+  indices = np.argsort(fvalues)
+  if reverse:
+    indices = np.flip(indices)
+
+  if start is not None or end is not None:
+    fvalues = fvalues[indices]
+    start_index = bisect.bisect(fvalues, start) if start is not None else 0
+    end_index = bisect.bisect(fvalues, end) if end is not None else len(indices)
+    if start_index > end_index:
+      start_index, end_index = end_index, start_index
+
+    indices = indices[start_index: end_index]
+
+  return indices
+
+
+class StreamSortedScan:
+
+  def __init__(self, reader, field,
+               start=None,
+               end=None,
+               slice_size=None,
+               max_slices=None,
+               reverse=False):
+    self._slice_size = slice_size or 100000
+    self._max_slices = max_slices or 16
+    self._reader = reader
+    self._slices = collections.OrderedDict()
+    self._indices = _compute_indices(reader, field, start=start, end=end, reverse=reverse)
+
+  def _get_slice(self, idx):
+    sidx = (idx // self._slice_size) * self._slice_size
+    data = self._slices.get(sidx)
+    if data is None:
+      if len(self._slices) >= self._max_slices:
+        self._slices.popitem(last=False)
+
+      slice_size = min(self._slice_size, len(self._indices) - sidx)
+
+      data = self._reader.get_slice(sidx, size=slice_size)
+      self._slices[sidx] = data
+    else:
+      self._slices.move_to_end(sidx)
+
+    return data, idx - sidx
+
+  def _as_numpy(self, rdata):
+    return {field: np.array(data) for field, data in rdata.items()}
+
+  def scan(self):
+    # An ampty array can contain fields which are Python lists, so _as_numpy() is
+    # used when returning data to the caller.
+    rdata = self._reader.empty_array(self._slice_size)
+    widx = 0
+    for idx in self._indices:
+      if widx == self._slice_size:
+        yield widx, self._as_numpy(rdata)
+        widx = 0
+
+      sdata, sidx = self._get_slice(idx)
+      for field, data in rdata.items():
+        data[widx] = sdata[field][sidx]
+
+      widx += 1
+
+    if widx:
+      frdata = collections.OrderedDict()
+      for field, data in rdata.items():
+        frdata[field] = data[: widx]
+
+      yield widx, self._as_numpy(frdata)
+

py_misc_utils/streamed_file.py

@@ -0,0 +1,144 @@
+import os
+import tempfile
+import threading
+
+from . import alog
+from . import assert_checks as tas
+from . import fin_wrap as fw
+
+
+class StreamedFile:
+
+  def __init__(self, resp):
+    self._resp = resp
+    self._lock = threading.Lock()
+    self._cond = threading.Condition(lock=self._lock)
+
+    tmpfile = tempfile.TemporaryFile()
+    fw.fin_wrap(self, '_tempfile', tmpfile, finfn=tmpfile.close)
+
+    self._offset = 0
+    self._size = 0
+    self._completed = False
+    self._closed = False
+    self._thread = threading.Thread(target=self._stream, daemon=True)
+    self._thread.start()
+
+  def _stream(self):
+    for data in self._resp:
+      with self._lock:
+        self._tempfile.seek(self._size)
+        self._tempfile.write(data)
+        self._size += len(data)
+        self._cond.notify_all()
+        if self._closed:
+          break
+
+    with self._lock:
+      self._completed = True
+      self._cond.notify_all()
+
+  def _wait_completed(self):
+    with self._lock:
+      while not (self._completed or self._closed):
+        self._cond.wait()
+
+  def close(self):
+    with self._lock:
+      self._closed = True
+      while not self._completed:
+        self._cond.wait()
+
+    self._thread.join()
+
+    with self._lock:
+      tempfile = self._tempfile
+      if tempfile is not None:
+        fw.fin_wrap(self, '_tempfile', None)
+
+    if tempfile is not None:
+      tempfile.close()
+
+  @property
+  def closed(self):
+    return self._tempfile is None
+
+  def seek(self, pos, whence=os.SEEK_SET):
+    if whence == os.SEEK_SET:
+      offset = pos
+    elif whence == os.SEEK_CUR:
+      offset = self._offset + pos
+    elif whence == os.SEEK_END:
+      self._wait_completed()
+      offset = self._size + pos
+    else:
+      alog.xraise(ValueError, f'Invalid seek mode: {whence}')
+
+    if offset > 0:
+      if whence != os.SEEK_END:
+        self._wait_completed()
+      tas.check_le(offset, self._size, msg=f'Offset out of range')
+
+    tas.check_ge(offset, 0, msg=f'Offset out of range')
+
+    self._offset = offset
+
+    return offset
+
+  def tell(self):
+    return self._offset
+
+  def _read(self, offset, size, adj_offset):
+    while not (self._completed or self._closed or
+               (size >= 0 and self._size >= offset + size)):
+      self._cond.wait()
+
+    available = self._size - offset
+    to_read = min(size, available) if size >= 0 else available
+    if not self._closed and to_read > 0:
+      self._tempfile.seek(offset)
+      data = self._tempfile.read(to_read)
+      if adj_offset:
+        self._offset += len(data)
+    else:
+      data = b''
+
+    return data
+
+  def read(self, size=-1):
+    with self._lock:
+      return self._read(self._offset, size, True)
+
+  def read1(self, size=-1):
+    return self.read(size=size)
+
+  def peek(self, size=0):
+    with self._lock:
+      size = min(size, max(1, self._size - self._offset))
+
+      return self._read(self._offset, size, False) if size > 0 else b''
+
+  def pread(self, offset, size):
+    with self._lock:
+      return self._read(offset, size, False)
+
+  def flush(self):
+    pass
+
+  def readable(self):
+    return not self.closed
+
+  def seekable(self):
+    return not self.closed
+
+  def writable(self):
+    return False
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, *exc):
+    self.close()
+
+    return False
+

py_misc_utils/tempdir.py

@@ -0,0 +1,79 @@
+import os
+import shutil
+import tempfile
+
+from . import cleanups
+from . import global_namespace as gns
+from . import rnd_utils as rngu
+
+
+class _RootDir:
+
+  def __init__(self):
+    self._path = tempfile.mkdtemp()
+    self._cid = cleanups.register(shutil.rmtree, self._path, ignore_errors=True)
+
+  def create(self):
+    return tempfile.mkdtemp(dir=self._path)
+
+  def root(self):
+    return self._path
+
+
+_ROOTDIR = gns.Var(f'{__name__}.ROOTDIR',
+                   fork_init=True,
+                   defval=lambda: _RootDir())
+
+def _root_dir():
+  return gns.get(_ROOTDIR)
+
+
+def create():
+  return _root_dir().create()
+
+
+def get_temp_root():
+  return _root_dir().root()
+
+
+def _try_fastfs_dir(path):
+  if os.path.isdir(path):
+    fastfs_dir = os.path.join(path, 'fastfs')
+    try:
+      os.makedirs(fastfs_dir, exist_ok=True)
+
+      return fastfs_dir
+    except:
+      pass
+
+
+def _find_fastfs_dir():
+  fastfs_dirs = []
+
+  if (path := os.getenv('FASTFS_DIR')) is not None:
+    fastfs_dirs.append(path)
+
+  if os.name == 'posix':
+    # Try known tmpfs/ramfs places in case on Linux.
+    fastfs_dirs.append(f'/run/user/{os.getuid()}')
+    fastfs_dirs.append('/dev/shm')
+
+  fastfs_dirs.append(tempfile.gettempdir())
+  fastfs_dirs.append(os.getcwd())
+
+  for path in fastfs_dirs:
+    if (fastfs_dir := _try_fastfs_dir(path)) is not None:
+      return fastfs_dir
+
+
+_FASTFS_DIR = _find_fastfs_dir()
+_NAMELEN = int(os.getenv('FASTFS_NAMELEN', 12))
+
+def fastfs_dir(name=None, namelen=_NAMELEN):
+  name = name or rngu.rand_string(namelen)
+
+  path = os.path.join(_FASTFS_DIR, name)
+  os.makedirs(path, exist_ok=True)
+
+  return path
+

py_misc_utils/template_replace.py

@@ -0,0 +1,51 @@
+import re
+import string
+
+from . import alog
+
+
+class _FnDict:
+
+  def __init__(self, lookup_fn):
+    self._lookup_fn = lookup_fn
+
+  def __getitem__(self, key):
+    m = re.match(r'([^:]+):(.*)', key)
+    if m:
+      lkey, defval = m.group(1), m.group(2)
+    else:
+      lkey, defval = key, None
+
+    return self._lookup_fn(lkey, defval=defval)
+
+
+def _dict_lookup_fn(vals, delim, misses_ok):
+
+  def lookup_fn(key, defval=None):
+    value = vals.get(key, defval)
+    if value is None:
+      if not misses_ok:
+        alog.xraise(KeyError, f'String template replace missing value for key: {key}')
+      else:
+        value = f'{delim}{key}'
+
+    return value
+
+  return lookup_fn
+
+
+def template_replace(st, vals=None, lookup_fn=None, delim=None, misses_ok=None):
+  delim = delim or '$'
+  misses_ok = False if misses_ok is None else misses_ok
+
+  class Template(string.Template):
+
+    # Allow for brace ID with the format ${ID:DEFAULT_VALUE}.
+    braceidpattern = r'((?a:[_a-z][_a-z0-9]*)(:[^}]*)?)'
+    delimiter = delim
+
+  if lookup_fn is None:
+    lookup_fn = _dict_lookup_fn(vals, delim, misses_ok)
+
+  return Template(st).safe_substitute(_FnDict(lookup_fn))
+