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

py_misc_utils/moving_average.py

@@ -0,0 +1,19 @@
+from . import assert_checks as tas
+from . import num_utils as nu
+
+
+class MovingAverage:
+
+  def __init__(self, factor, init=None):
+    tas.check(factor >= 0.0 and factor <= 1.0, msg=f'Invalid factor: {factor:.4e}')
+    self._factor = factor
+    self.value = init
+
+  def update(self, value):
+    if self.value is None:
+      self.value = value
+    else:
+      self.value = nu.mix(self.value, value, self._factor)
+
+    return self.value
+

py_misc_utils/msgpack_streamer.py

@@ -0,0 +1,26 @@
+import msgpack
+
+from . import alog
+from . import gfs
+from . import utils as ut
+
+
+class MsgPackStreamer:
+
+  def __init__(self, url, rename_columns=None, **kwargs):
+    self._url = url
+    self._rename_columns = ut.value_or(rename_columns, dict())
+    self._kwargs = kwargs
+
+  def generate(self):
+    with gfs.open(self._url, mode='rb', **self._kwargs) as stream:
+      unpacker = msgpack.Unpacker(stream)
+      for recd in unpacker:
+        for key, name in self._rename_columns.items():
+          recd[name] = recd.pop(key)
+
+        yield recd
+
+  def __iter__(self):
+    return self.generate()
+

py_misc_utils/multi_wait.py

@@ -0,0 +1,24 @@
+import threading
+
+
+class MultiWait:
+
+  def __init__(self, count):
+    self._count = count
+    self._sigcount = 0
+    self._lock = threading.Lock()
+    self._cond = threading.Condition(lock=self._lock)
+
+  def signal(self, n=1):
+    with self._lock:
+      self._sigcount = min(self._sigcount + n, self._count)
+      if self._sigcount == self._count:
+        self._cond.notify_all()
+
+  def wait(self, timeout=None):
+    with self._lock:
+      if self._count > self._sigcount:
+        return self._cond.wait(timeout=timeout)
+
+    return True
+

py_misc_utils/multiprocessing.py

@@ -0,0 +1,102 @@
+import functools
+import gc
+import multiprocessing
+import signal
+import sys
+
+from . import alog
+from . import cleanups
+from . import global_namespace as gns
+
+
+class TerminationError(Exception):
+  pass
+
+def _sig_handler(sig, frame):
+  if sig == signal.SIGINT:
+    raise KeyboardInterrupt()
+  elif sig == signal.SIGTERM:
+    raise TerminationError()
+
+
+def _process_setup():
+  signal.signal(signal.SIGINT, _sig_handler)
+  signal.signal(signal.SIGTERM, _sig_handler)
+
+
+def _process_cleanup():
+  cleanups.run()
+  gc.collect()
+  # Ignore {INT, TERM} signals on exit.
+  signal.signal(signal.SIGINT, signal.SIG_IGN)
+  signal.signal(signal.SIGTERM, signal.SIG_IGN)
+
+
+_GNS_KEY = 'gns'
+
+def _wrap_procfn_parent(method):
+  ctx = dict(method=method)
+  ctx.update({_GNS_KEY: gns.parent_switch(method)})
+
+  return ctx
+
+
+def _wrap_procfn_child(method, pctx):
+  parent_gns = pctx.pop(_GNS_KEY, None)
+  if parent_gns is not None:
+    gns.child_switch(method, parent_gns)
+
+  return pctx
+
+
+_CONTEXT_KEY = '_parent_context'
+
+def _capture_parent_context(method, kwargs):
+  pctx = _wrap_procfn_parent(method)
+  kwargs.update({_CONTEXT_KEY: pctx})
+
+  return kwargs
+
+
+def _apply_child_context(kwargs):
+  pctx = kwargs.pop(_CONTEXT_KEY)
+  pctx = _wrap_procfn_child(pctx['method'], pctx)
+
+  return kwargs
+
+
+def procfn_wrap(procfn, *args, **kwargs):
+  try:
+    _process_setup()
+
+    return procfn(*args, **kwargs)
+  except (KeyboardInterrupt, TerminationError):
+    sys.exit(1)
+  except Exception as ex:
+    alog.exception(ex, exmsg=f'Exception while running process function')
+    raise
+  finally:
+    _process_cleanup()
+
+
+def _wrapped_procfn(procfn, *args, **kwargs):
+  kwargs = _apply_child_context(kwargs)
+
+  return procfn(*args, **kwargs)
+
+
+def create_process(procfn, args=(), kwargs=None, context=None, daemon=None):
+  if context is None:
+    mpctx = multiprocessing
+  elif isinstance(context, str):
+    mpctx = multiprocessing.get_context(method=context)
+  else:
+    mpctx = context
+
+  kwargs = dict() if kwargs is None else kwargs.copy()
+
+  kwargs = _capture_parent_context(mpctx.get_start_method(), kwargs)
+  target = functools.partial(procfn_wrap, _wrapped_procfn, procfn, *args, **kwargs)
+
+  return mpctx.Process(target=target, daemon=daemon)
+

py_misc_utils/named_array.py

@@ -0,0 +1,224 @@
+import array
+import re
+
+import numpy as np
+import pandas as pd
+
+from . import assert_checks as tas
+from . import core_utils as cu
+from . import np_utils as npu
+
+
+_NOT_NUMERIC = 'xS'
+
+
+class Field:
+
+  __slots__ = ('name', 'data', 'size', 'fmt', 'str_tbl')
+
+  def __init__(self, name, data, size, fmt, str_tbl):
+    self.name = name
+    self.data = data
+    self.size = size
+    self.fmt = fmt
+    self.str_tbl = str_tbl
+
+  def np_array(self):
+    arr = np.array(self.data)
+
+    return arr.reshape(-1, self.size) if self.size > 1 else arr
+
+  def append(self, arg):
+    if self.fmt == 'S':
+      arg = self.stringify(arg)
+
+    if self.size == 1:
+      self.data.append(arg)
+    else:
+      assert self.size == len(arg), f'{self.name}({self.size}) vs. {len(arg)}'
+      self.data.extend(arg)
+
+  def extend(self, arg):
+    assert len(arg) % self.size == 0, f'{self.name}({self.size}) vs. {len(arg)}'
+    if self.fmt == 'S':
+      arg = self.stringify(arg)
+
+    if isinstance(arg, np.ndarray) and isinstance(self.data, array.array):
+      nptype = np.dtype(self.data.typecode)
+      if nptype != arg.dtype:
+        arg = arg.astype(nptype)
+      self.data.frombytes(arg.tobytes())
+    else:
+      self.data.extend(arg)
+
+  def stringify(self, arg):
+    if self.size == 1:
+      return self.str_tbl.add(arg)
+
+    return tuple(self.str_tbl.add(x) for x in arg)
+
+  def __len__(self):
+    return len(self.data) // self.size
+
+  def __getitem__(self, i):
+    if self.size == 1:
+      return self.data[i]
+    elif isinstance(i, int):
+      offset = i * self.size
+
+      return self.data[offset: offset + self.size]
+    else:
+      start, stop, step = (x * self.size for x in i.indices(len(self)))
+      data = []
+      for n in range(start, stop, step):
+        data.append(self.data[n: n + self.size])
+
+      return data
+
+  @staticmethod
+  def create(name, fmt, str_tbl):
+    m = re.match(r'(\d+)([a-zA-Z])', fmt)
+    if m:
+      size, efmt = int(m.group(1)), m.group(2)
+    else:
+      size, efmt = 1, fmt
+
+    data = array.array(efmt) if not efmt in _NOT_NUMERIC else []
+
+    return Field(name, data, size, efmt, str_tbl)
+
+
+class NamedArray:
+
+  def __init__(self, names, fmt=None):
+    fnames = re.split(r'\s*,\s*', names) if isinstance(names, str) else names
+    if fmt is None:
+      cnames, ffmt = [], []
+      for name_format in fnames:
+        m = re.match(r'(\w+)\s*=\s*(\w+)$', name_format)
+        tas.check(m, msg=f'Invalid name=format specification: {name_format}')
+        cnames.append(m.group(1))
+        ffmt.append(m.group(2))
+
+      fnames = cnames
+    else:
+      ffmt = re.findall(r'\d*[a-zA-Z]', fmt) if isinstance(fmt, str) else fmt
+      tas.check_eq(len(fnames), len(ffmt),
+                   msg=f'Mismatching names and format sizes: {fnames} vs {ffmt}')
+
+    str_tbl = cu.StringTable()
+    fields = dict()
+    for name, fmt in zip(fnames, ffmt):
+      fields[name] = Field.create(name, fmt, str_tbl)
+
+    self._fields = fields
+    self._fieldseq = tuple(fields.values())
+    self._str_tbl = str_tbl
+
+  def add_column(self, name, fmt, data):
+    tas.check(name not in self._fields, msg=f'Column "{name}" already exists')
+
+    field = Field.create(name, fmt, self._str_tbl)
+
+    if isinstance(data, np.ndarray):
+      data = data.flatten()
+
+    field.extend(data)
+
+    self._fields[name] = field
+    self._fieldseq += (field,)
+
+  def append(self, *args):
+    for field, arg in zip(self._fieldseq, args):
+      field.append(arg)
+
+  def extend(self, other):
+    for field, ofield in zip(self._fieldseq, other._fieldseq):
+      field.extend(ofield.data)
+
+  def append_extend(self, *args):
+    for field, arg in zip(self._fieldseq, args):
+      field.extend(arg)
+
+  def kwappend(self, **kwargs):
+    for name, field in self._fields.items():
+      arg = kwargs[name]
+      field.append(arg)
+
+  def get_tuple_item(self, i):
+    item = []
+    for field in self._fieldseq:
+      item.append(field[i])
+
+    return tuple(item)
+
+  def get_arrays(self, names=None):
+    if names:
+      return tuple(self._fields[name].np_array() for name in names)
+
+    return tuple(field.np_array() for field in self._fieldseq)
+
+  def get_array(self, name):
+    return self._fields[name].np_array()
+
+  def data(self):
+    return {field.name: field.np_array() for field in self._fieldseq}
+
+  def to_numpy(self, dtype=None):
+    for field in self._fieldseq:
+      tas.check(field.fmt not in _NOT_NUMERIC,
+                msg=f'Only purely numeric arrays can be converted to ' \
+                f'numpy: {field.name}={field.fmt}')
+      tas.check_eq(field.size, 1,
+                   msg=f'Only fields with size==1 can be converted to ' \
+                   f'numpy: {field.name}={field.size}')
+
+    if dtype is None:
+      dtype = npu.infer_np_dtype(self.dtypes())
+
+    na = np.empty(self.shape, dtype=dtype)
+    for i, field in enumerate(self._fieldseq):
+      na[:, i] = field.data
+
+    return na
+
+  def __len__(self):
+    size = None
+    for field in self._fieldseq:
+      fsize = len(field)
+      if size is None:
+        size = fsize
+      else:
+        tas.check_eq(fsize, size,
+                     msg=f'Unmatching size for "{field.name}": {fsize} vs. {size}')
+
+    return size or 0
+
+  def __getitem__(self, i):
+    item = dict()
+    for field in self._fieldseq:
+      item[field.name] = field[i]
+
+    return item
+
+  def __array__(self, dtype=None):
+    return self.to_numpy(dtype=dtype)
+
+  def to_dataframe(self):
+    df_data = dict()
+    for name, arr in self.data().items():
+      df_data[name] = arr if arr.ndim == 1 else arr.tolist()
+
+    return pd.DataFrame(data=df_data)
+
+  @property
+  def shape(self):
+    return (len(self), len(self._fieldseq))
+
+  def dtypes(self):
+    types = []
+    for field in self._fieldseq:
+      types.append(np.dtype('O') if field.fmt in _NOT_NUMERIC else np.dtype(field.fmt))
+
+    return tuple(types)
+

py_misc_utils/no_break.py

@@ -0,0 +1,46 @@
+import signal
+
+from . import core_utils as cu
+from . import signal as sgn
+
+
+class NoBreak:
+
+  SIGMAP = {
+    'INT': signal.SIGINT,
+    'TERM': signal.SIGTERM,
+  }
+
+  def __init__(self, sigs=None, exit_trigger=False):
+    if sigs is None:
+      self._signals = tuple(self.SIGMAP.values())
+    else:
+      if isinstance(sigs, str):
+        sigs = cu.splitstrip(sigs, ',')
+
+      self._signals = tuple(self.SIGMAP[sig] for sig in sigs)
+
+    self._exit_trigger = exit_trigger
+
+  def __enter__(self):
+    self._signal_received = []
+    for sig in self._signals:
+      sgn.signal(sig, self._handler, prio=sgn.STD_PRIO)
+
+    return self
+
+  def _handler(self, sig, frame):
+    self._signal_received.append((sig, frame))
+
+    return sgn.HANDLED
+
+  def __exit__(self, *exc):
+    for sig in self._signals:
+      sgn.unsignal(sig, self._handler)
+
+    if self._exit_trigger:
+      for sig, frame in self._signal_received:
+        sgn.trigger(sig, frame)
+
+    return False
+

py_misc_utils/no_except.py

@@ -0,0 +1,32 @@
+# Using logging module directly here to avoid import the alog module, since this
+# is supposed to be a lower level module with minimal/no local dependencies.
+import logging
+import traceback
+
+
+def no_except(fn, *args, **kwargs):
+  try:
+    return fn(*args, **kwargs)
+  except Exception as ex:
+    # The logging.exception() API does emit in DEBUG level, and here we want something
+    # a bit more higher level.
+    tb = traceback.format_exc()
+    logging.warning(f'Exception while running function: {ex}\n{tb}')
+
+    return ex
+
+
+def qno_except(fn, *args, **kwargs):
+  try:
+    return fn(*args, **kwargs)
+  except Exception as ex:
+    return ex
+
+
+def xwrap_fn(fn, *args, **kwargs):
+
+  def fwrap():
+    return no_except(fn, *args, **kwargs)
+
+  return fwrap
+

py_misc_utils/np_ml_framework.py

@@ -0,0 +1,184 @@
+import collections
+import os
+
+
+_MODULES = dict()
+
+
+def _register_attr(mod, name, attr):
+  xa = getattr(mod, name, None)
+  if xa is not None:
+    raise RuntimeError(f'Attribute "{name}" already exists: {xa}')
+
+  setattr(mod, name, attr)
+
+
+def _register(name, mod, checkfn, fromfn, attrs):
+  _MODULES[name] = mod
+  mod.__npml_name = name
+  mod.__npml_check = checkfn
+  mod.__npml_from = fromfn
+
+  for attr_name, attr_value in attrs.items():
+    _register_attr(mod, attr_name, attr_value)
+
+
+def _parse_priorities():
+  prefs = os.getenv('NPML_ORDER', 'torch,np,jax,tf').split(',')
+
+  return {mod: len(prefs) - i for i, mod in enumerate(prefs)}
+
+
+# Numpy
+try:
+  import numpy as np
+
+  def _np_from(mod, t, tref):
+    if mod is not None:
+      if mod is torch or mod is tfnp:
+        return t.numpy()
+
+    return np.asarray(t)
+
+  def _np_check(t):
+    return isinstance(t, np.ndarray)
+
+  _register('np', np, _np_check, _np_from,
+            {
+              'item': lambda t: t.item(),
+              'tolist': lambda t: t.tolist(),
+            })
+except ImportError:
+  np = None
+
+
+# PyTorch
+try:
+  import torch
+  from torch.utils import dlpack as torch_dlpack
+
+  def _torch_from(mod, t, tref):
+    if mod is not None:
+      if mod is np:
+        return torch.from_numpy(t).to(tref.device)
+      if mod is jaxnp:
+        return torch_dlpack.from_dlpack(jax_dlpack.to_dlpack(t)).to(tref.device)
+      if mod is tfnp:
+        return torch_dlpack.from_dlpack(tf_dlpack.to_dlpack(t)).to(tref.device)
+
+    return torch.tensor(t).to(tref.device)
+
+  def _torch_check(t):
+    return isinstance(t, torch.Tensor)
+
+  _register('torch', torch, _torch_check, _torch_from,
+            {
+              'item': lambda t: t.item(),
+              'tolist': lambda t: t.tolist(),
+            })
+except ImportError:
+  torch = None
+
+
+# JAX
+try:
+  import jax
+  from jax import dlpack as jax_dlpack
+  import jax.numpy as jaxnp
+
+  def _jaxdev(t):
+    return next(iter(t.devices()))
+
+  def _jax_from(mod, t, tref):
+    if mod is not None:
+      if mod is torch:
+        return jax.device_put(jax_dlpack.from_dlpack(torch_dlpack.to_dlpack(t)), _jaxdev(tref))
+      if mod is tfnp:
+        return jax.device_put(jax_dlpack.from_dlpack(tf_dlpack.to_dlpack(t)), _jaxdev(tref))
+
+    return jax.device_put(jaxnp.asarray(t), _jaxdev(tref))
+
+  def _jax_check(t):
+    return isinstance(t, jax.Array)
+
+  _register('jax', jaxnp, _jax_check, _jax_from,
+            {
+              'item': lambda t: t.item(),
+              'tolist': lambda t: t.tolist(),
+            })
+except ImportError:
+  jaxnp = None
+
+
+# Tensorflow
+try:
+  import tensorflow as tf
+  import tensorflow.experimental.dlpack as tf_dlpack
+  import tensorflow.experimental.numpy as tfnp
+
+  def _tf_from(mod, t, tref):
+    if mod is not None:
+      if mod is torch:
+        with tf.device(tref.device):
+          return tf_dlpack.from_dlpack(torch_dlpack.to_dlpack(t))
+      if mod is jaxnp:
+        with tf.device(tref.device):
+          return tf_dlpack.from_dlpack(jax_dlpack.to_dlpack(t))
+
+    with tf.device(tref.device):
+      return tf.convert_to_tensor(t)
+
+  def _tf_check(t):
+    return tf.is_tensor(t)
+
+  tfnp.experimental_enable_numpy_behavior()
+  _register('tf', tfnp, _tf_check, _tf_from,
+            {
+              'item': lambda t: t.item(),
+              'tolist': lambda t: t.tolist(),
+            })
+except ImportError:
+  tfnp = None
+
+
+_DEFAULT_MODULE = os.getenv('NPML_DEFAULT', 'np')
+
+if _DEFAULT_MODULE not in _MODULES:
+  raise RuntimeError(f'Unable to find default Numpy ML module: {_DEFAULT_MODULE}')
+
+
+_MODULES_PRIORITY = _parse_priorities()
+_MODULES_SEQ = tuple(_MODULES[x] for x in sorted(list(_MODULES.keys()),
+                                                 key=lambda m: _MODULES_PRIORITY.get(m, -1),
+                                                 reverse=True))
+
+
+def _get_module(t):
+  for mod in _MODULES_SEQ:
+    if mod.__npml_check(t):
+      return mod
+
+
+def resolve(*args):
+  mods = []
+  tprio, tmod, tref = -1, None, None
+  for i, t in enumerate(args):
+    mod = _get_module(t)
+    mods.append(mod)
+    if mod is not None:
+      prio = _MODULES_PRIORITY[mod.__npml_name]
+      if prio > tprio:
+        tmod = mod
+        tprio = prio
+        tref = t
+
+  if tmod is None:
+    return _MODULES[_DEFAULT_MODULE], args
+
+  rargs = list(args)
+  for i, mod in enumerate(mods):
+    if tmod is not mod:
+      rargs[i] = tmod.__npml_from(mod, args[i], tref)
+
+  return tmod, tuple(rargs)
+