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

py_misc_utils/ordered_set.py

@@ -0,0 +1,71 @@
+import copy
+
+
+class OrderedSet:
+
+  def __init__(self, init=None):
+    self._data = dict()
+    self._seqno = 0
+
+    for value in init or ():
+      self.add(value)
+
+  def add(self, value):
+    n = self._data.get(value)
+    if n is None:
+      self._data[value] = n = self._seqno
+      self._seqno += 1
+
+    return n
+
+  def remove(self, value):
+    self._data.pop(value)
+
+  def discard(self, value):
+    self._data.pop(value, None)
+
+  def pop(self):
+    value, n = self._data.popitem()
+
+    return value
+
+  def clear(self):
+    self._data = dict()
+    self._seqno = 0
+
+  def __len__(self):
+    return len(self._data)
+
+  def values(self):
+    return (y[0] for y in sorted(self._data.items(), key=lambda x: x[1]))
+
+  def __iter__(self):
+    return iter(self.values())
+
+  def __contains__(self, value):
+    return value in self._data
+
+  def union(self, *others):
+    nos = copy.copy(self)
+    for other in others:
+      for value in other:
+        nos.add(value)
+
+    return nos
+
+  def intersection(self, *others):
+    nos = OrderedSet()
+    for value in self.values():
+      if all(value in other for other in others):
+        nos.add(value)
+
+    return nos
+
+  def difference(self, *others):
+    nos = OrderedSet()
+    for value in self.values():
+      if not any(value in other for other in others):
+        nos.add(value)
+
+    return nos
+

py_misc_utils/osfd.py

@@ -0,0 +1,27 @@
+import os
+
+
+class OsFd:
+
+  def __init__(self, path, flags, remove_on_error=None, **kwargs):
+    self._path = path
+    self._flags = flags
+    self._kwargs = kwargs
+    self._remove_on_error = (remove_on_error if remove_on_error is not None else
+                             not os.path.exists(path))
+
+  def __enter__(self):
+    self._fd = os.open(self._path, self._flags, **self._kwargs)
+
+    return self._fd
+
+  def __exit__(self, *exc):
+    os.close(self._fd)
+    if any(ex is not None for ex in exc) and self._remove_on_error:
+      try:
+        os.remove(self._path)
+      except OSError:
+        pass
+
+    return False
+

py_misc_utils/packet.py

@@ -0,0 +1,22 @@
+import os
+import struct
+
+from . import iostream as ios
+
+
+_SIZE_PACKER = struct.Struct('<Q')
+
+def write_packet(fd, data):
+  packet = _SIZE_PACKER.pack(len(data)) + data
+  iofd = ios.IOStream(fd)
+  iofd.write(packet)
+
+
+def read_packet(fd):
+  iofd = ios.IOStream(fd)
+  data = iofd.read(_SIZE_PACKER.size)
+  size = _SIZE_PACKER.unpack(data)[0]
+  packet = iofd.read(size)
+
+  return packet
+

py_misc_utils/parquet_streamer.py

@@ -0,0 +1,69 @@
+import contextlib
+import functools
+
+import pyarrow.parquet as pq
+
+from . import alog
+from . import fin_wrap as fw
+from . import gfs
+from . import url_fetcher as urlf
+from . import utils as ut
+
+
+class ParquetStreamer:
+
+  def __init__(self, url,
+               batch_size=128,
+               load_columns=None,
+               rename_columns=None,
+               num_workers=None,
+               **kwargs):
+    self._url = url
+    self._batch_size = batch_size
+    self._load_columns = ut.value_or(load_columns, dict())
+    self._rename_columns = ut.value_or(rename_columns, dict())
+    self._num_workers = num_workers
+    self._kwargs = kwargs
+
+  def _fetcher(self):
+    if self._load_columns:
+      return urlf.UrlFetcher(num_workers=self._num_workers,
+                             fs_kwargs=self._kwargs)
+    else:
+      return contextlib.nullcontext()
+
+  def _prefetch(self, fetcher, recs):
+    if fetcher is not None:
+      for recd in recs:
+        for key in self._load_columns.keys():
+          fetcher.enqueue(recd[key])
+
+  def _transform(self, fetcher, recd):
+    if fetcher is not None:
+      for key, name in self._load_columns.items():
+        recd[name] = fetcher.wait(recd[key])
+
+    for key, name in self._rename_columns.items():
+      recd[name] = recd.pop(key)
+
+    return recd
+
+  def generate(self):
+    with (self._fetcher() as fetcher,
+          gfs.open(self._url, mode='rb', **self._kwargs) as stream):
+      pqfd = pq.ParquetFile(stream)
+      for batch in pqfd.iter_batches(batch_size=self._batch_size):
+        recs = batch.to_pylist()
+
+        self._prefetch(fetcher, recs)
+        for recd in recs:
+          try:
+            yield self._transform(fetcher, recd)
+          except GeneratorExit:
+            raise
+          except Exception as ex:
+            alog.verbose(f'Unable to create parquet entry ({recd}): {ex}')
+
+  def __iter__(self):
+    return self.generate()
+

py_misc_utils/pd_utils.py

@@ -0,0 +1,254 @@
+import array
+import collections
+import datetime
+import os
+import re
+
+import numpy as np
+import pandas as pd
+
+from . import alog
+from . import assert_checks as tas
+from . import core_utils as cu
+from . import gfs
+from . import np_utils as npu
+from . import utils as ut
+
+
+def get_df_columns(df, discards=None):
+  dset = discards or {}
+
+  return [c for c in df.columns if c not in dset]
+
+
+def get_typed_columns(df, type_fn, discards=None):
+  cols = []
+  for c in get_df_columns(df, discards=discards):
+    if type_fn(df[c].dtype):
+      cols.append(c)
+
+  return cols
+
+
+def re_select_columns(df, re_cols):
+  cols = []
+  for c in df.columns:
+    for rc in re_cols:
+      if re.match(rc, c):
+        cols.append(c)
+        break
+
+  return cols
+
+
+def read_csv(path, rows_sample=100, dtype=None, args=None):
+  with gfs.open(path, mode='r') as fd:
+    args = dict() if args is None else args
+    if args.get('index_col') is None:
+      args = args.copy()
+      fields = ut.comma_split(fd.readline())
+      fd.seek(0)
+      # If 'index_col' is not specified, we use column 0 if its name is empty, otherwise
+      # we disable it setting it to False.
+      args['index_col'] = False if fields[0] else 0
+
+    if dtype is None:
+      return pd.read_csv(fd, **args)
+    if cu.isdict(dtype):
+      dtype = {c: np.dtype(t) for c, t in dtype.items()}
+    else:
+      df_test = pd.read_csv(fd, nrows=rows_sample, **args)
+      fd.seek(0)
+      dtype = {c: dtype for c in get_typed_columns(df_test, npu.is_numeric)}
+
+    return pd.read_csv(fd, dtype=dtype, **args)
+
+
+def save_dataframe(df, path, **kwargs):
+  _, ext = os.path.splitext(os.path.basename(path))
+  if ext == '.pkl':
+    args = ut.dict_subset(kwargs, ('compression', 'protocol', 'storage_options'))
+    if 'protocol' not in args:
+      args['protocol'] = ut.pickle_proto()
+    with gfs.open(path, mode='wb') as fd:
+      df.to_pickle(fd, **args)
+  elif ext == '.csv':
+    args = ut.dict_subset(kwargs, ('float_format', 'columns', 'header', 'index',
+                                   'index_label', 'mode', 'encoding', 'quoting',
+                                   'quotechar', 'line_terminator', 'chunksize',
+                                   'date_format', 'doublequote', 'escapechar',
+                                   'decimal', 'compression', 'error', 'storage_options'))
+
+    # For CSV file, unless otherwise specified, and the index has no name, drop
+    # the index column as it adds no value to the output (it's simply a sequential).
+    if not df.index.name:
+      args = ut.dict_setmissing(args, index=None)
+
+    with gfs.open(path, mode='w') as fd:
+      df.to_csv(fd, **args)
+  else:
+    alog.xraise(RuntimeError, f'Unknown extension: {ext}')
+
+
+def load_dataframe(path, **kwargs):
+  _, ext = os.path.splitext(os.path.basename(path))
+  if ext == '.pkl':
+    with gfs.open(path, mode='rb') as fd:
+      return pd.read_pickle(fd)
+  elif ext == '.csv':
+    rows_sample = kwargs.pop('rows_sample', 100)
+    dtype = kwargs.pop('dtype', None)
+    args = ut.dict_subset(kwargs, ('sep', 'delimiter', 'header', 'names',
+                                   'index_col', 'usecols', 'squeeze', 'prefix',
+                                   'mangle_dupe_cols', 'dtype', 'engine',
+                                   'converters', 'true_values', 'false_values',
+                                   'skipinitialspace', 'skiprows', 'skipfooter',
+                                   'nrows', 'na_values', 'keep_default_na', 'na_filter',
+                                   'verbose', 'skip_blank_lines', 'parse_dates',
+                                   'infer_datetime_format', 'keep_date_col', 'date_parser',
+                                   'dayfirst', 'cache_dates', 'iterator', 'chunksize',
+                                   'compression', 'thousands', 'decimal', 'lineterminator',
+                                   'quotechar', 'quoting', 'doublequote', 'escapechar',
+                                   'comment', 'encoding', 'dialect', 'error_bad_lines',
+                                   'warn_bad_lines', 'delim_whitespace', 'low_memory',
+                                   'memory_map', 'float_precision', 'storage_options'))
+
+    return read_csv(path, rows_sample=rows_sample, dtype=dtype,
+                    args=args)
+  else:
+    alog.xraise(RuntimeError, f'Unknown extension: {ext}')
+
+
+def to_npdict(df, reset_index=False, dtype=None, no_convert=()):
+  if reset_index and df.index.name:
+    df = df.reset_index()
+
+  cdata = dict()
+  for c in df.columns:
+    data = df[c].to_numpy()
+    if dtype is not None and c not in no_convert:
+      data = npu.astype(data, c, dtype)
+    cdata[c] = data
+
+  return cdata
+
+
+def load_dataframe_as_npdict(path, reset_index=False, dtype=None, no_convert=()):
+  df = load_dataframe(path)
+
+  return to_npdict(df, reset_index=reset_index, dtype=dtype, no_convert=no_convert)
+
+
+def column_or_index(df, name, numpy=True):
+  data = df.get(name)
+  if data is None and df.index.name == name:
+    data = df.index
+  if data is not None:
+    return data.to_numpy() if numpy else data
+
+
+def columns_transform(df, cols, tfn):
+  for c in cols:
+    cv = df.get(c)
+    if cv is not None:
+      df[c] = tfn(c, cv, index=False)
+    elif df.index.name == c:
+      df.index = pd.Index(data=tfn(c, df.index, index=True), name=df.index.name)
+    else:
+      alog.xraise(RuntimeError, f'Unable to find column or index named "{c}"')
+
+  return df
+
+
+def get_columns_index(df):
+  cols = df.columns.tolist()
+
+  return ut.make_index_dict(cols), cols
+
+
+def concat_dataframes(files, **kwargs):
+  dfs = []
+  for path in files:
+    df = load_dataframe(path, **kwargs)
+    dfs.append(df)
+
+  return pd.concat(dfs, **kwargs) if dfs else None
+
+
+def get_dataframe_groups(df, cols, cols_transforms=None):
+  # Pandas groupby() is extremely slow when there are many groups as it builds
+  # a DataFrame for each group. This simply collects the rows associated with each
+  # tuple values representing the grouped columns.
+  # Row numbers must be strictly ascending within each group, do NOT change that!
+  groups = collections.defaultdict(lambda: array.array('L'))
+  if cols_transforms:
+    tcols = [(df[c], cols_transforms.get(c, pycu.ident)) for c in cols]
+    for i in range(len(df)):
+      k = tuple([f(d[i]) for d, f in tcols])
+      groups[k].append(i)
+  else:
+    cdata = [df[c] for c in cols]
+    for i in range(len(df)):
+      k = tuple([d[i] for d in cdata])
+      groups[k].append(i)
+
+  return groups
+
+
+def limit_per_group(df, cols, limit):
+  mask = np.full(len(df), False)
+
+  groups = get_dataframe_groups(df, cols)
+  for k, g in groups.items():
+    if limit > 0:
+      rows = g[: limit]
+    else:
+      rows = g[limit:]
+    mask[rows] = True
+
+  return df[mask]
+
+
+def dataframe_column_rewrite(df, name, fn):
+  data = df.get(name)
+  if data is not None:
+    df[name] = fn(data.to_numpy())
+  elif df.index.name == name:
+    nvalues = fn(df.index.to_numpy())
+    df.index = type(df.index)(data=nvalues, name=df.index.name)
+  else:
+    alog.xraise(RuntimeError, f'No column or index named "{name}"')
+
+
+def correlate(df, col, top_n=None):
+  ccorr = df.corrwith(df[col])
+  scorr = ccorr.sort_values(key=lambda x: abs(x), ascending=False)
+  top_n = len(scorr) if top_n is None else min(top_n, len(scorr))
+
+  return tuple([(scorr.index[i], scorr[i]) for i in range(top_n)])
+
+
+def type_convert_dataframe(df, types):
+  # Pandas DataFrame astype() doe not have an 'inplace' argument.
+  for c, t in types.items():
+    df[c] = df[c].astype(t)
+
+  return df
+
+
+def dataframe_rows_select(df, indices):
+  return df.loc[df.index[indices]]
+
+
+def sorted_index(df, col):
+  if df.index.name == col:
+    return df.sort_values(col)
+
+  sdf = df.sort_values(col, ignore_index=True)
+
+  return sdf.set_index(col)
+
+
+def datetime_to_epoch(data):
+  return (pd.to_datetime(data) - datetime.datetime(1970, 1, 1)).dt.total_seconds()
+

py_misc_utils/periodic_task.py

@@ -0,0 +1,61 @@
+import threading
+import weakref
+
+from . import alog
+from . import scheduler as sch
+from . import weak_call as wcall
+
+
+class PeriodicTask:
+
+  def __init__(self, name, periodic_fn, period, scheduler=None, stop_on_error=None):
+    self._name = name
+    self._periodic_fn = wcall.WeakCall(periodic_fn)
+    self._period = period
+    self._scheduler = scheduler or sch.common_scheduler()
+    self._stop_on_error = stop_on_error in (None, True)
+    self._lock = threading.Lock()
+    self._event = None
+    self._completed_event = None
+
+  def start(self):
+    with self._lock:
+      if self._event is None:
+        self._schedule()
+
+    return self
+
+  def stop(self):
+    completed_event = None
+    with self._lock:
+      if self._event is not None:
+        # If "events" is empty, we were not able to cancel the task, so it will be
+        # in flight, and we need to wait for it to complete before exiting.
+        events = self._scheduler.cancel(self._event)
+        completed_event = self._completed_event if not events else None
+        self._event = None
+
+    if completed_event is not None:
+      completed_event.wait()
+
+  def _schedule(self):
+    self._event = self._scheduler.enter(self._period, self._runner)
+    self._completed_event = threading.Event()
+
+  def _runner(self):
+    re_issue = True
+    try:
+      if self._periodic_fn() is wcall.GONE:
+        re_issue = False
+    except Exception as ex:
+      alog.exception(ex, exmsg=f'Exception while running periodic task "{self._name}"')
+      re_issue = not self._stop_on_error
+    finally:
+      with self._lock:
+        self._completed_event.set()
+        if self._event is not None:
+          if re_issue:
+            self._schedule()
+          else:
+            self._event = None
+

py_misc_utils/pickle_wrap.py

@@ -0,0 +1,121 @@
+import importlib
+import os
+import pickle
+import sys
+
+from . import alog
+from . import core_utils as cu
+from . import inspect_utils as iu
+from . import std_module as stdm
+
+
+KNOWN_MODULES = {
+  'builtins',
+  'numpy',
+  'pandas',
+  'torch',
+  cu.root_module(__name__),
+}
+
+def add_known_module(modname):
+  KNOWN_MODULES.add(cu.root_module(modname))
+
+
+def _needs_wrap(obj):
+  objmod = iu.moduleof(obj)
+  if objmod is not None:
+    modname = cu.root_module(objmod)
+    if modname in KNOWN_MODULES or stdm.is_std_module(modname):
+      return False
+
+  return True
+
+
+def _wrap(obj, pickle_module):
+  wrapped = 0
+  if isinstance(obj, (list, tuple)):
+    wobj = []
+    for v in obj:
+      wv = _wrap(v, pickle_module)
+      if wv is not v:
+        wrapped += 1
+
+      wobj.append(wv)
+
+    return type(obj)(wobj) if wrapped else obj
+  elif cu.isdict(obj):
+    wobj = type(obj)()
+    for k, v in obj.items():
+      wk = _wrap(k, pickle_module)
+      if wk is not k:
+        wrapped += 1
+      wv = _wrap(v, pickle_module)
+      if wv is not v:
+        wrapped += 1
+
+      wobj[wk] = wv
+
+    return wobj if wrapped else obj
+  elif isinstance(obj, PickleWrap):
+    return obj
+  elif _needs_wrap(obj):
+    return PickleWrap(obj, pickle_module=pickle_module)
+  else:
+    return obj
+
+
+def _unwrap(obj, pickle_module):
+  unwrapped = 0
+  if isinstance(obj, (list, tuple)):
+    uwobj = []
+    for v in obj:
+      wv = _unwrap(v, pickle_module)
+      if wv is not v:
+        unwrapped += 1
+
+      uwobj.append(wv)
+
+    return type(obj)(uwobj) if unwrapped else obj
+  elif cu.isdict(obj):
+    uwobj = type(obj)()
+    for k, v in obj.items():
+      wk = _unwrap(k, pickle_module)
+      if wk is not k:
+        unwrapped += 1
+      wv = _unwrap(v, pickle_module)
+      if wv is not v:
+        unwrapped += 1
+
+      uwobj[wk] = wv
+
+    return uwobj if unwrapped else obj
+  elif isinstance(obj, PickleWrap):
+    try:
+      return obj.load(pickle_module=pickle_module)
+    except Exception as ex:
+      alog.debug(f'Unable to reload pickle-wrapped data ({obj.wrapped_class()}): {ex}')
+      return obj
+  else:
+    return obj
+
+
+def wrap(obj, pickle_module=pickle):
+  return _wrap(obj, pickle_module)
+
+
+def unwrap(obj, pickle_module=pickle):
+  return _unwrap(obj, pickle_module)
+
+
+class PickleWrap:
+
+  def __init__(self, obj, pickle_module=pickle):
+    self._class = iu.qual_name(obj)
+    self._data = pickle_module.dumps(obj)
+
+  def wrapped_class(self):
+    return self._class
+
+  def load(self, pickle_module=pickle):
+    return pickle_module.loads(self._data)
+