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

py_misc_utils/abs_timeout.py

@@ -0,0 +1,12 @@
+import time
+
+
+class AbsTimeout:
+
+  def __init__(self, timeout, timefn=None):
+    self._timefn = time.time if timefn is None else timefn
+    self._expires = self._timefn() + timeout if timeout is not None else None
+
+  def get(self):
+    return max(0, self._expires - self._timefn()) if self._expires is not None else None
+

py_misc_utils/alog.py

@@ -0,0 +1,311 @@
+import argparse
+import logging
+import math
+import os
+import sys
+import time
+import traceback
+import types
+
+from . import call_limiter as cl
+from . import run_once as ro
+from . import traceback as tb
+
+
+DEBUG = logging.DEBUG
+INFO = logging.INFO
+WARNING = logging.WARNING
+ERROR = logging.ERROR
+CRITICAL = logging.CRITICAL
+
+SPAM = DEBUG - 2
+VERBOSE = DEBUG - 1
+DEBUG0 = DEBUG + 1
+DEBUG1 = DEBUG + 2
+DEBUG2 = DEBUG + 3
+DEBUG3 = DEBUG + 4
+
+_SHORT_LEV = {
+  SPAM: 'SP',
+  VERBOSE: 'VB',
+  DEBUG0: '0D',
+  DEBUG1: '1D',
+  DEBUG2: '2D',
+  DEBUG3: '3D',
+  DEBUG: 'DD',
+  INFO: 'IN',
+  WARNING: 'WA',
+  ERROR: 'ER',
+  CRITICAL: 'CR',
+}
+
+
+class Formatter(logging.Formatter):
+
+  def __init__(self, emit_extra=None):
+    super().__init__()
+    self.emit_extra = emit_extra
+
+  def format(self, r):
+    hdr = self.make_header(r)
+    msg = (r.msg % r.args) if r.args else r.msg
+
+    return '\n'.join([f'{hdr}: {ln}' for ln in msg.split('\n')])
+
+  def formatTime(self, r, datefmt=None):
+    if datefmt:
+      return time.strftime(datefmt, r.created)
+
+    tstr = time.strftime('%Y%m%d %H:%M:%S', time.localtime(r.created))
+    usecs = math.modf(r.created)[0] * 1e6
+
+    return f'{tstr}.{usecs:06.0f}'
+
+  def make_header(self, r):
+    tstr = self.formatTime(r)
+    lid = _SHORT_LEV.get(r.levelno, r.levelname[:2])
+    hdr = f'{lid}{tstr};{os.getpid()};{r.module}'
+    if self.emit_extra:
+      extras = [str(getattr(r, name, None)) for name in self.emit_extra]
+      hdr = f'{hdr};{";".join(extras)}'
+
+    return hdr
+
+
+_DEFAULT_ARGS = dict(
+  log_level=os.getenv('LOG_LEVEL', 'INFO'),
+  log_file=os.getenv('LOG_FILE', 'STDERR'),
+  log_mod_levels=[],
+  log_emit_extra=[],
+)
+
+def add_logging_options(parser):
+  parser.add_argument('--log_level', type=str, default=_DEFAULT_ARGS.get('log_level'),
+                      choices={'SPAM', 'VERBOSE', 'DEBUG', 'DEBUG0', 'DEBUG1', 'DEBUG2',
+                               'DEBUG3', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'},
+                      help='The logging level')
+  parser.add_argument('--log_file', type=str, default=_DEFAULT_ARGS.get('log_file'),
+                      help='Comma separated list of target log files (STDOUT, STDERR ' \
+                      f'are also recognized)')
+  parser.add_argument('--log_mod_levels', nargs='*',
+                      help='Comma separated list of LOGGER_NAME,LEVEL to set the log level at')
+  parser.add_argument('--log_emit_extra', nargs='*',
+                      help='Which other logging record fields should be emitted')
+
+
+@ro.run_once
+def _add_levels():
+  logging.addLevelName(SPAM, 'SPAM')
+  logging.addLevelName(VERBOSE, 'VERBOSE')
+  logging.addLevelName(DEBUG0, 'DEBUG0')
+  logging.addLevelName(DEBUG1, 'DEBUG1')
+  logging.addLevelName(DEBUG2, 'DEBUG2')
+  logging.addLevelName(DEBUG3, 'DEBUG3')
+
+
+def _clear_logging_handlers():
+  # Python >= 3.8 has a force=True argument to logging.basicConfig() to force
+  # initialization, but since Colab is not there yet, we do it manually.
+  root_logger = logging.getLogger()
+  for handler in tuple(root_logger.handlers):
+    handler.flush()
+    root_logger.removeHandler(handler)
+
+
+def _set_logmod_levels(mlevels):
+  mlevels = list(mlevels) if mlevels else []
+  env_mlevels = os.getenv('LOGMOD_LEVELS', None)
+  if env_mlevels is not None:
+    mlevels.extend(env_mlevels.split(':'))
+  for mlev in mlevels:
+    mod, level = mlev.split(',')
+    logging.getLogger(mod).setLevel(logging.getLevelName(level.upper()))
+
+
+def setup_logging(args):
+  _add_levels()
+  _clear_logging_handlers()
+
+  numeric_level = logging.getLevelName(args.log_level.upper())
+  handlers = []
+  if args.log_file:
+    for fname in args.log_file.split(','):
+      if fname == 'STDOUT':
+        handler = logging.StreamHandler(sys.stdout)
+      elif fname == 'STDERR':
+        handler = logging.StreamHandler(sys.stderr)
+      else:
+        handler = logging.StreamHandler(open(fname, mode='a'))
+
+      handler.setLevel(numeric_level)
+      handler.setFormatter(Formatter(emit_extra=args.log_emit_extra))
+      handlers.append(handler)
+
+  logging.basicConfig(level=numeric_level, handlers=handlers, force=True)
+
+  set_current_level(numeric_level, set_logger=False)
+
+  _set_logmod_levels(args.log_mod_levels)
+
+
+def basic_setup(**kwargs):
+  args = _DEFAULT_ARGS.copy()
+  args.update(kwargs)
+  setup_logging(types.SimpleNamespace(**args))
+
+
+def get_main_config():
+  return types.SimpleNamespace(add_arguments=add_logging_options,
+                               config_module=setup_logging)
+
+
+_LEVEL = DEBUG
+
+def set_current_level(level, set_logger=True):
+  if set_logger:
+    logger = logging.getLogger()
+    logger.setLevel(level)
+    for handler in logger.handlers:
+      handler.setLevel(level)
+
+  global _LEVEL
+
+  _LEVEL = level
+
+
+def level_active(level):
+  return _LEVEL <= level
+
+
+def level_run(level, fn):
+  return fn() if level_active(level) else None
+
+
+_LOGGING_FRAMES = 1 if sys.version_info >= (3, 11) else 2
+
+def logging_args(kwargs):
+  limit = kwargs.pop('limit', -1)
+  stacklevel = kwargs.get('stacklevel', 1)
+  if limit < 0 or cl.trigger(__file__, limit):
+    kwargs['stacklevel'] = stacklevel + _LOGGING_FRAMES
+
+    return kwargs
+
+
+def _nested_args(kwargs):
+  kwargs['stacklevel'] = kwargs.get('stacklevel', 1) + 1
+
+  return kwargs
+
+
+def _dmsg(msg):
+  return msg() if callable(msg) else msg
+
+
+def log(level, msg, *args, **kwargs):
+  kwargs = logging_args(kwargs)
+  if kwargs is not None:
+    logging.log(level, _dmsg(msg), *args, **kwargs)
+
+
+def spam(msg, *args, **kwargs):
+  if SPAM >= _LEVEL:
+    log(SPAM, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def verbose(msg, *args, **kwargs):
+  if VERBOSE >= _LEVEL:
+    log(VERBOSE, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def debug0(msg, *args, **kwargs):
+  if DEBUG0 >= _LEVEL:
+    log(DEBUG0, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def debug1(msg, *args, **kwargs):
+  if DEBUG1 >= _LEVEL:
+    log(DEBUG1, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def debug2(msg, *args, **kwargs):
+  if DEBUG2 >= _LEVEL:
+    log(DEBUG2, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def debug3(msg, *args, **kwargs):
+  if DEBUG3 >= _LEVEL:
+    log(DEBUG3, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def debug(msg, *args, **kwargs):
+  if DEBUG >= _LEVEL:
+    log(DEBUG, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def info(msg, *args, **kwargs):
+  if INFO >= _LEVEL:
+    log(INFO, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def warning(msg, *args, **kwargs):
+  if WARNING >= _LEVEL:
+    log(WARNING, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def error(msg, *args, **kwargs):
+  if ERROR >= _LEVEL:
+    log(ERROR, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def critical(msg, *args, **kwargs):
+  if CRITICAL >= _LEVEL:
+    log(CRITICAL, _dmsg(msg), *args, **_nested_args(kwargs))
+
+
+def exception(e, *args, **kwargs):
+  kwargs = logging_args(kwargs)
+  if kwargs is not None:
+    msg = kwargs.pop('exmsg', 'Exception')
+    tb = traceback.format_exc()
+    error(f'{_dmsg(msg)}: {e}\n{tb}', *args, **_nested_args(kwargs))
+
+
+def xraise(e, msg, *args, **kwargs):
+  if kwargs.pop('logit', False):
+    error(msg, *args, **_nested_args(kwargs))
+
+  raise e(_dmsg(msg))
+
+
+def async_log(level, msg, *args, **kwargs):
+  # This one cannot use the logging module as it could be called from signal
+  # handler asycnhronously. The logging.getLevelName() is safe since it is simply
+  # a (lockless) dictionary lookup.
+  # Similarly, no other APIs taking locks can be called from this context.
+  if level >= _LEVEL:
+    kwargs = logging_args(kwargs)
+    if kwargs is not None:
+      # Fake a logging record, filling up only the fields used by the Formatter.
+      # Do not call logging APIs for that, for the same reasons cited above.
+      frame = tb.get_frame(n=1)
+      module = frame.f_globals.get('__name__', 'ASYNC').split('.')[-1]
+
+      now = time.time()
+      record = types.SimpleNamespace(
+        msg=_dmsg(msg),
+        args=args,
+        created=now,
+        msecs=math.modf(now)[0] * 1000,
+        levelno=level,
+        levelname=logging.getLevelName(level),
+        module=module,
+      )
+
+      formatter = Formatter()
+
+      logfd = kwargs.pop('file', sys.stderr)
+      logfd.write(formatter.format(record))
+      logfd.write('\n')
+      logfd.flush()
+

py_misc_utils/app_main.py

@@ -0,0 +1,179 @@
+import argparse
+import functools
+import inspect
+import sys
+import typing
+import yaml
+
+from . import alog
+from . import core_utils as cu
+from . import global_namespace as gns
+from . import multiprocessing as mp
+
+
+def _get_init_modules():
+  # Here is the place to import (import here to avoid cycling dependencies) and
+  # call the get_main_config() API of modules which require setting up a
+  # command line and configuring themselves with the parsed arguments.
+  # Note that alog is imported at the top since it is used in other places (and
+  # also has minimal dependencies which do not create issues).
+  # Objects returned by the get_main_config() API must have a add_arguments(parser)
+  # API to allow them to add command line arguments, and a config_module(args) API
+  # to configure themselves with the parsed arguments.
+  # Example:
+  #
+  #   from . import foo
+  #   modules.append(foo.get_main_config())
+  #
+  modules = []
+  modules.append(alog.get_main_config())
+
+  return tuple(modules)
+
+
+def _add_arguments(init_modules, parser):
+  for module in init_modules:
+    module.add_arguments(parser)
+
+
+def _config_modules(init_modules, args):
+  for module in init_modules:
+    module.config_module(args)
+
+
+def _child_setup_modules(args):
+  init_modules = _get_init_modules()
+  _config_modules(init_modules, args)
+
+  return args
+
+
+_ARGS = gns.Var(f'{__name__}.ARGS', child_fn=_child_setup_modules)
+
+def _main(parser, mainfn, args, rem_args):
+  if isinstance(mainfn, Main):
+    mainfn.add_arguments(parser)
+
+  init_modules = _get_init_modules()
+  _add_arguments(init_modules, parser)
+
+  if rem_args:
+    xargs = args or sys.argv[1:]
+
+    ddpos = cu.lindex(xargs, '--')
+    if ddpos >= 0:
+      rargs = xargs[ddpos + 1:]
+      xargs = xargs[: ddpos]
+    else:
+      rargs = []
+
+    parsed_args = parser.parse_args(args=xargs)
+    setattr(parsed_args, rem_args, tuple(rargs))
+  else:
+    parsed_args = parser.parse_args(args=args)
+
+  gns.set(_ARGS, parsed_args)
+  _config_modules(init_modules, parsed_args)
+
+  mainfn(parsed_args)
+
+
+def main(parser, mainfn, args=None, rem_args=None):
+  mp.procfn_wrap(_main, parser, mainfn, args, rem_args)
+
+
+def basic_main(mainfn, description='Basic Main'):
+  parser = argparse.ArgumentParser(
+    description=description,
+    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+  )
+  main(parser, mainfn)
+
+
+def _child_setup_functions(setup_functions):
+  for setupfn in setup_functions:
+    setupfn()
+
+  return setup_functions
+
+
+_SETUP_FUNCTIONS = gns.Var(f'{__name__}.SETUP_FUNCTIONS',
+                           child_fn=_child_setup_functions,
+                           defval=[])
+
+def add_setupfn(setupfn, run=True):
+  if run:
+    setupfn()
+
+  setup_functions = gns.get(_SETUP_FUNCTIONS)
+  setup_functions.append(setupfn)
+
+
+# This is similar to Fire but brings up the app_main infrastructure.
+# Use as:
+#
+# @app_main.Main
+# def my_main(arg, ..., kwarg=17, ...):
+#   ...
+#
+# if __name__ == '__main__':
+#   parser = argparse.ArgumentParser(...)
+#   ...
+#   app_main.main(parser, my_main, ...)
+#
+class Main:
+
+  def __init__(self, func):
+    self._func = func
+    self._sig = inspect.signature(func)
+    functools.update_wrapper(self, func)
+
+  def __call__(self, parsed_args):
+    args, kwargs = [], {}
+    for n, p in self._sig.parameters.items():
+      pv = getattr(parsed_args, n, None)
+      if p.kind == p.POSITIONAL_ONLY:
+        args.append(pv)
+      else:
+        kwargs[n] = pv
+
+    return self._func(*args, **kwargs)
+
+  def add_arguments(self, parser):
+    fname = self._func.__name__
+
+    for n, p in self._sig.parameters.items():
+      choices = None
+      defval = p.default if p.default is not p.empty else None
+      if p.annotation is not p.empty:
+        ptype = p.annotation
+        if typing.get_origin(ptype) == typing.Literal:
+          choices = typing.get_args(ptype)
+          ptype = type(choices[0])
+
+        type_cast = functools.partial(cu.to_type, vtype=ptype)
+      elif defval is not None:
+        ptype = type(defval)
+        type_cast = functools.partial(cu.to_type, vtype=ptype)
+      else:
+        ptype, type_cast = str, yaml.safe_load
+
+      action = argparse.BooleanOptionalAction if ptype is bool else None
+
+      help_str = f'Argument "{n}" (type={ptype.__name__}) of function {fname}(...)'
+      if p.default is p.empty or p.kind == p.POSITIONAL_ONLY:
+        parser.add_argument(n,
+                            metavar=n.upper(),
+                            action=action,
+                            type=type_cast,
+                            default=defval,
+                            choices=choices,
+                            help=help_str)
+      else:
+        parser.add_argument(f'--{n}',
+                            action=action,
+                            type=type_cast,
+                            default=defval,
+                            choices=choices,
+                            help=help_str)
+

py_misc_utils/archive_streamer.py

@@ -0,0 +1,112 @@
+import collections
+import hashlib
+import os
+import tarfile
+import zipfile
+
+from . import alog
+from . import assert_checks as tas
+from . import gfs
+from . import img_utils as imgu
+from . import utils as ut
+
+
+ArchiveSpecs = collections.namedtuple('ArchiveSpecs', 'kind, compression, base_path, purl')
+ArchiveEntry = collections.namedtuple('ArchiveEntry', 'name, data')
+
+
+_EXT_COMPRESSION = {
+  'gz': 'gz',
+  'xz': 'xz',
+  'bz2': 'bz2',
+  'bzip2': 'bz2',
+}
+
+def parse_specs(url):
+  usplit = gfs.splitext(url)
+
+  compression = _EXT_COMPRESSION.get(usplit.ext)
+  ubase = usplit.base if compression else usplit.purl.path
+
+  base_path, ext = os.path.splitext(ubase)
+
+  tas.check(ext, msg=f'Unable to infer archive type: {url}')
+
+  return ArchiveSpecs(kind=usplit.ext.lower(),
+                      compression=compression,
+                      base_path=base_path,
+                      purl=usplit.purl)
+
+
+class ArchiveStreamer:
+
+  def __init__(self, url, **kwargs):
+    self._url = url
+    self._kwargs = kwargs
+
+  def _url_uid(self, url):
+    return hashlib.sha1(url.encode()).hexdigest()[: 8]
+
+  def _generate_zip(self, specs):
+    # The ZIP format requires random access (specifically, the file list is at EOF)
+    # so it is better to cache the file locally before opening.
+    with gfs.open_local(self._url, mode='rb', **self._kwargs) as stream:
+      zfile = zipfile.ZipFile(stream, mode='r')
+      for zinfo in zfile.infolist():
+        if not zinfo.is_dir():
+          data = zfile.read(zinfo)
+          yield ArchiveEntry(name=zinfo.filename, data=data)
+
+  def _generate_tar(self, specs):
+    with gfs.open(self._url, mode='rb', **self._kwargs) as stream:
+      tfile = tarfile.open(mode=f'r|{specs.compression or ""}', fileobj=stream)
+      for tinfo in tfile:
+        data = tfile.extractfile(tinfo).read()
+        yield ArchiveEntry(name=tinfo.name, data=data)
+
+  def _generate_parquet(self, specs):
+    # Keep the import dependency local, to make it required only if parquet is used.
+    from . import parquet_streamer as pqs
+
+    uid = self._url_uid(self._url)
+
+    pq_streamer = pqs.ParquetStreamer(self._url, **self._kwargs)
+    for i, recd in enumerate(pq_streamer):
+      # Simulate a streaming similar to what a Web Dataset would expect, with a
+      # UID.ENTITY naming, where the UID is constant for all the entities of a record
+      # (which are streamed sequentially).
+      ruid = f'{uid}_{i}'
+      for name, data in recd.items():
+        yield ArchiveEntry(name=f'{ruid}.{name}', data=data)
+
+  def _generate_msgpack(self, specs):
+    # Keep the import dependency local, to make it required only if parquet is used.
+    from . import msgpack_streamer as mps
+
+    uid = self._url_uid(self._url)
+
+    mps_streamer = mps.MsgPackStreamer(self._url, **self._kwargs)
+    for i, recd in enumerate(mps_streamer):
+      # Simulate a streaming similar to what a Web Dataset would expect, with a
+      # UID.ENTITY naming, where the UID is constant for all the entities of a record
+      # (which are streamed sequentially).
+      ruid = f'{uid}_{i}'
+      for name, data in recd.items():
+        yield ArchiveEntry(name=f'{ruid}.{name}', data=data)
+
+  def generate(self):
+    specs = parse_specs(self._url)
+    if specs.kind == 'zip':
+      yield from self._generate_zip(specs)
+    elif specs.kind == 'tar':
+      yield from self._generate_tar(specs)
+    elif specs.kind == 'parquet':
+      yield from self._generate_parquet(specs)
+    elif specs.kind == 'msgpack':
+      yield from self._generate_msgpack(specs)
+    else:
+      alog.xraise(RuntimeError, f'Unknown archive type "{specs.kind}": {self._url}')
+
+  def __iter__(self):
+    return self.generate()
+