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

py_misc_utils/assert_checks.py

@@ -0,0 +1,118 @@
+import inspect
+import logging
+import os
+
+# Note: This is a core module, which cannot have local imports!
+# Using the logging module above to not add the alog module dependency here,
+# since that requires more dependencies to be pulled.
+# At the end, alog.log(...) is exactly the same as logging.log(...) in substance.
+
+
+def _get_loc(path, lineno):
+  if os.path.isfile(path):
+    with open(path, mode='r') as fd:
+      lines = fd.read().splitlines()
+
+    return lines[lineno - 1] if len(lines) > lineno else None
+
+
+def _get_caller_info(n_back):
+  frame = inspect.stack()[n_back + 1][0]
+  caller = inspect.getframeinfo(frame)
+  loc = _get_loc(caller.filename, caller.lineno)
+  if loc:
+    return f'{caller.filename}:{caller.lineno}: {loc.lstrip()}'
+
+  return f'{caller.filename}:{caller.lineno}'
+
+
+def _report_fail(level, op, *args, **kwargs):
+  fmsg = kwargs.get('msg')
+  cinfo = _get_caller_info(2)
+  if fmsg:
+    cinfo = f'{cinfo}; {fmsg}'
+  if op:
+    if callable(op):
+      fname = getattr(op, '__name__', str(op))
+      arg_list = ', '.join([str(arg) for arg in args])
+      res = kwargs.get('res')
+      if res:
+        res_op = kwargs['res_op']
+        msg = f'{fname}({arg_list}) {res_op} {res} failed from {cinfo}'
+      else:
+        msg = f'{fname}({arg_list}) failed from {cinfo}'
+    else:
+      assert len(args) == 2, len(args)
+      msg = f'{args[0]} {op} {args[1]} failed from {cinfo}'
+  else:
+    msg = f'Check failed from {cinfo}'
+
+  logging.log(level, msg)
+
+  raise AssertionError(msg)
+
+
+def check(a, level=logging.ERROR, msg=None):
+  if not a:
+    _report_fail(level, None, msg=msg)
+
+
+def check_fn(fn, *args, level=logging.ERROR, msg=None):
+  if not fn(*args):
+    _report_fail(level, fn, *args, msg=msg)
+
+
+def check_fnres_eq(res, fn, *args, level=logging.ERROR, msg=None):
+  if not (fn(*args) == res):
+    _report_fail(level, fn, *args, res=res, res_op='==', msg=msg)
+
+
+def check_fnres_ne(res, fn, *args, level=logging.ERROR, msg=None):
+  if fn(*args) == res:
+    _report_fail(level, fn, *args, res=res, res_op='!=', msg=msg)
+
+
+def check_is_none(a, level=logging.ERROR, msg=None):
+  if a is not None:
+    _report_fail(level, '==', a, 'None', msg=msg)
+
+
+def check_is_not_none(a, level=logging.ERROR, msg=None):
+  if a is None:
+    _report_fail(level, '!=', a, 'None', msg=msg)
+
+
+def check_in(a, b, level=logging.ERROR, msg=None):
+  if a not in b:
+    _report_fail(level, 'in', a, b, msg=msg)
+
+
+def check_eq(a, b, level=logging.ERROR, msg=None):
+  if not (a == b):
+    _report_fail(level, '==', a, b, msg=msg)
+
+
+def check_ne(a, b, level=logging.ERROR, msg=None):
+  if not (a != b):
+    _report_fail(level, '!=', a, b, msg=msg)
+
+
+def check_le(a, b, level=logging.ERROR, msg=None):
+  if not (a <= b):
+    _report_fail(level, '<=', a, b, msg=msg)
+
+
+def check_ge(a, b, level=logging.ERROR, msg=None):
+  if not (a >= b):
+    _report_fail(level, '>=', a, b, msg=msg)
+
+
+def check_lt(a, b, level=logging.ERROR, msg=None):
+  if not (a < b):
+    _report_fail(level, '<', a, b, msg=msg)
+
+
+def check_gt(a, b, level=logging.ERROR, msg=None):
+  if not (a > b):
+    _report_fail(level, '>', a, b, msg=msg)
+

py_misc_utils/ast_utils.py

@@ -0,0 +1,121 @@
+import ast
+import logging
+import os
+
+from . import alog
+
+
+def _ends_with_return(slist):
+  return slist and isinstance(slist[-1], ast.Return)
+
+
+def _ifize_stmt_list(slist):
+  ni = None
+  for i, node in enumerate(slist):
+    if isinstance(node, ast.If):
+      ni = i
+      break
+
+  if ni is not None and ni + 1 < len(slist):
+    ifnode = slist[ni]
+    orelse = ifnode.orelse or []
+
+    # If one branch of an IF ends with RETURN, and the other does not,
+    # move the remaining of the statements after the IF, within the branch
+    # which does not have the RETURN.
+    # Example (open 'else'), turns:
+    #
+    #   op1
+    #   if cond:
+    #     op2
+    #     return x
+    #   else:
+    #     op3
+    #   op4
+    #   return y
+    #
+    # Into:
+    #   op1
+    #   if cond:
+    #     op2
+    #     return x
+    #   else:
+    #     op3
+    #     op4
+    #     return y
+    #
+    # Example (open 'if'), turns:
+    #
+    #   op1
+    #   if cond:
+    #     op2
+    #   else:
+    #     op3
+    #     return x
+    #   op4
+    #   return y
+    #
+    # Into:
+    #   op1
+    #   if cond:
+    #     op2
+    #     op4
+    #     return y
+    #   else:
+    #     op3
+    #     return x
+    #
+    # IOW, it makes sure that if within an AST instruction list, the first IF (if any),
+    # will be the last instruction of the list.
+    if _ends_with_return(ifnode.body):
+      if not _ends_with_return(orelse):
+        remlist = _ifize_stmt_list(slist[ni + 1:])
+        orelse.extend(remlist)
+        ifnode.orelse = orelse
+        slist = slist[: ni + 1]
+
+    elif _ends_with_return(orelse):
+      remlist = _ifize_stmt_list(slist[ni + 1:])
+      ifnode.body.extend(remlist)
+      slist = slist[: ni + 1]
+
+  return slist
+
+
+def ifize(node):
+  for field, value in ast.iter_fields(node):
+    if isinstance(value, list):
+      for lvalue in value:
+        if isinstance(lvalue, ast.AST):
+          ifize(lvalue)
+
+      xlist = _ifize_stmt_list(value)
+      if xlist is not value:
+        setattr(node, field, xlist)
+    elif isinstance(value, ast.AST):
+      ifize(value)
+
+
+def dump(node, indent=None):
+  if indent is None:
+    indent = os.getenv('AST_INDENT')
+    if indent is not None:
+      indent = int(indent)
+
+  return ast.dump(node, indent=indent)
+
+
+def static_eval(node, eval_globals, eval_locals, filename=None):
+  if isinstance(node, ast.stmt):
+    mod = ast.Module(body=[node], type_ignores=[])
+    cmod = compile(mod, filename=filename or '<ast_eval>', mode='exec')
+    exec(cmod, eval_globals, eval_locals)
+  elif isinstance(node, ast.expr):
+    expr = ast.Expression(body=node)
+    cexpr = compile(expr, filename=filename or '<ast_eval>', mode='eval')
+    value = eval(cexpr, eval_globals, eval_locals)
+
+    return value
+  else:
+    alog.xraise(ValueError, f'Invalid AST node: {dump(node)}')
+

py_misc_utils/async_manager.py

@@ -0,0 +1,189 @@
+import asyncio
+import collections
+import multiprocessing
+import os
+import queue
+import threading
+
+import numpy as np
+
+from . import cleanups
+from . import global_namespace as gns
+from . import multiprocessing as mp
+from . import work_results as wres
+
+
+class AsyncContext:
+
+  def __init__(self):
+    self._contexts = dict()
+
+  async def add(self, name, context):
+    result = await context.__aenter__()
+    self._contexts[name] = (context, result)
+
+    return result
+
+  async def remove(self, name):
+    context, _ = self._contexts.pop(name)
+    await context.__aexit__(None, None, None)
+
+  async def get(self, name, context_ctor):
+    context_result = self._contexts.get(name)
+    if context_result is None:
+      result = await self.add(name, context_ctor())
+    else:
+      result = context_result[1]
+
+    return result
+
+  async def close(self, *exc):
+    rexc = exc or (None, None, None)
+    needs_raise = False
+    for context, _ in self._contexts.values():
+      exitres = await context.__aexit__(*rexc)
+      needs_raise = needs_raise or not exitres
+
+    self._contexts = dict()
+    if needs_raise and rexc[1] is not None:
+      raise rexc[1]
+
+  async def __aenter__(self):
+    return self
+
+  async def __aexit__(self, *exc):
+    await self.close(*exc)
+
+    return False
+
+
+Work = collections.namedtuple('Work', 'id, ctor')
+
+class _Worker:
+
+  def __init__(self, mpctx, wid, out_queue):
+    self._wid = wid
+    self._out_queue = out_queue
+    self._in_queue = mpctx.Queue()
+    self._proc = mp.create_process(self._run, context=mpctx)
+    self._proc.start()
+
+  def _run(self):
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    thread = threading.Thread(target=self._work_feeder, args=(loop,), daemon=True)
+    thread.start()
+
+    loop.run_forever()
+    thread.join()
+
+  async def _task_runner(self, context, work):
+    try:
+      task = work.ctor(context=context)
+
+      result = await task
+    except Exception as ex:
+      result = wres.WorkException(ex, workid=work.id)
+
+    self._out_queue.put((self._wid, work.id, result))
+
+  def _work_feeder(self, loop):
+    context = AsyncContext()
+
+    while True:
+      work = self._in_queue.get()
+      if work is None:
+        break
+
+      asyncio.run_coroutine_threadsafe(self._task_runner(context, work), loop)
+
+    asyncio.run_coroutine_threadsafe(self._shutdown(context, loop), loop)
+
+  @classmethod
+  async def _shutdown(cls, context, loop):
+    await context.close()
+    loop.stop()
+
+  def stop(self):
+    self._in_queue.put(None)
+    self._proc.join()
+
+  def enqueue_work(self, work_id, work_ctor):
+    self._in_queue.put(Work(id=work_id, ctor=work_ctor))
+
+
+class AsyncManager:
+
+  def __init__(self, num_workers=None, mpctx=multiprocessing):
+    num_workers = num_workers or os.cpu_count()
+
+    self._out_queue = mpctx.Queue()
+    self._workers = [_Worker(mpctx, i, self._out_queue) for i in range(num_workers)]
+    self._lock = threading.Lock()
+    self._queued = np.zeros(num_workers, dtype=np.int64)
+
+  def close(self):
+    for worker in self._workers:
+      worker.stop()
+
+  def enqueue_work(self, work_id, work_ctor):
+    with self._lock:
+      wid = np.argmin(self._queued)
+      self._queued[wid] += 1
+      worker = self._workers[wid]
+
+    worker.enqueue_work(work_id, work_ctor)
+
+  def fetch_result(self, block=True, timeout=None):
+    try:
+      wid, work_id, result = self._out_queue.get(block=block, timeout=timeout)
+
+      with self._lock:
+        self._queued[wid] -= 1
+
+      return work_id, result
+    except queue.Empty:
+      pass
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, *exc):
+    self.close()
+
+    return False
+
+
+class AsyncRunner:
+
+  def __init__(self):
+    self._loop = asyncio.new_event_loop()
+    self._thread = threading.Thread(target=self._async_runner, daemon=True)
+    self._thread.start()
+    self._cid = cleanups.register(self.stop)
+
+  def _async_runner(self):
+    asyncio.set_event_loop(self._loop)
+    self._loop.run_forever()
+
+  def stop(self):
+    self._loop.call_soon_threadsafe(self._loop.stop)
+    self._thread.join()
+
+  def run(self, coro):
+    return asyncio.run_coroutine_threadsafe(coro, self._loop)
+
+
+
+_ASYNC_RUNNER = gns.Var(f'{__name__}.ASYNC_RUNNER',
+                        fork_init=True,
+                        defval=lambda: AsyncRunner())
+
+def _async_runner():
+  return gns.get(_ASYNC_RUNNER)
+
+
+def run_async(coro):
+  return _async_runner().run(coro)
+

py_misc_utils/break_control.py

@@ -0,0 +1,63 @@
+import signal
+import threading
+
+from . import fin_wrap as fw
+from . import signal as sgn
+
+
+_LOCK = threading.Lock()
+_HANDLERS = set()
+
+
+def _handler(sig, frame):
+  with _LOCK:
+    for h in _HANDLERS:
+      h.trigger(frame)
+
+  return sgn.HANDLED
+
+
+class BreakControl:
+
+  def __init__(self):
+    self._hit = False
+    self._frame = None
+
+  def open(self):
+    with _LOCK:
+      if not _HANDLERS:
+        sgn.signal(signal.SIGINT, _handler)
+      _HANDLERS.add(self)
+
+    return self
+
+  def close(self):
+    with _LOCK:
+      _HANDLERS.remove(self)
+      if not _HANDLERS:
+        sgn.unsignal(signal.SIGINT, _handler)
+
+  def __enter__(self):
+    return self.open()
+
+  def __exit__(self, *exc):
+    self.close()
+
+    return False
+
+  def trigger(self, frame):
+    self._hit = True
+    self._frame = frame
+
+  def hit(self):
+    return self._hit
+
+  def frame(self):
+    return self._frame
+
+
+def create():
+  bc = BreakControl()
+
+  return fw.fin_wrap_np(bc.open(), bc.close)
+

py_misc_utils/buffered_iterator.py

@@ -0,0 +1,35 @@
+import collections
+
+
+IterData = collections.namedtuple('IterData', 'n, left, data')
+
+
+class BufferedIterator:
+
+  def __init__(self, data, buffer_size):
+    self._data = data
+    self._buffer_size = buffer_size
+
+  def generate(self):
+    queue, n = [], 0
+    for data in self._data:
+      if len(queue) < self._buffer_size:
+        queue.append(data)
+      else:
+        cidx = n % len(queue)
+        cdata = queue[cidx]
+        queue[cidx] = data
+        yield IterData(n=n, left=len(queue), data=cdata)
+
+      n += 1
+
+    for i in range(len(queue)):
+      cidx = n % len(queue)
+      cdata = queue[cidx]
+      yield IterData(n=n, left=len(queue) - i - 1, data=cdata)
+      n += 1
+
+  def __iter__(self):
+    return self.generate()
+
+