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

py_misc_utils/daemon_process.py

@@ -0,0 +1,252 @@
+import atexit
+import collections
+import functools
+import logging
+import multiprocessing
+import os
+import pickle
+import psutil
+import signal
+import sys
+import tempfile
+import time
+
+from . import fs_utils as fsu
+from . import lockfile as lockf
+from . import osfd
+from . import packet as pkt
+
+
+DaemonResult = collections.namedtuple(
+  'DaemonResult',
+  'pid, msg, exclass',
+  defaults=(-1, None, None),
+)
+
+def _get_pids_dir():
+  pidsdir = os.path.join(tempfile.gettempdir(), '.pids')
+  os.makedirs(pidsdir, exist_ok=True)
+
+  return pidsdir
+
+
+_PIDS_DIR = _get_pids_dir()
+
+def _get_pidfile(name):
+  return os.path.join(_PIDS_DIR, f'{name}.pid')
+
+
+def _term_handler(sig, frame):
+  sys.exit(sig)
+
+
+class DaemonBase:
+
+  def __init__(self, name):
+    self._name = name
+    self._pidfile = _get_pidfile(name)
+
+  def _write_result(self, wpipe, **kwargs):
+    dres = DaemonResult(**kwargs)
+    pkt.write_packet(wpipe, pickle.dumps(dres))
+
+  def _read_result(self, rpipe):
+    return pickle.loads(pkt.read_packet(rpipe))
+
+  def _delpid(self, pid=None):
+    with self._lockfile():
+      if pid is None or (xpid := self._readpid()) == pid:
+        try:
+          os.remove(self._pidfile)
+
+          return True
+        except OSError:
+          pass
+
+      return False
+
+  def _writepid(self, pid):
+    with self._lockfile():
+      xpid = self._readpid()
+      if xpid is not None:
+        if self._runnning_pid(xpid):
+          raise FileExistsError(f'Daemon already running with PID {xpid}')
+
+        os.remove(self._pidfile)
+
+      # Use mode=0o660 to make sure only allowed users can access the PID file.
+      with osfd.OsFd(self._pidfile, os.O_WRONLY | os.O_CREAT | os.O_EXCL, mode=0o660) as fd:
+        os.write(fd, f'{pid}\n'.encode())
+
+  def _readpid(self):
+    try:
+      with osfd.OsFd(self._pidfile, os.O_RDONLY) as fd:
+        return int(fsu.readall(fd).strip())
+    except IOError:
+      pass
+
+  def getpid(self):
+    with self._lockfile():
+      pid = self._readpid()
+
+      return pid if pid is None or self._runnning_pid(pid) else None
+
+  def _lockfile(self):
+    return lockf.LockFile(self._pidfile)
+
+  def _runnning_pid(self, pid):
+    try:
+      proc = psutil.Process(pid)
+
+      return proc.status() not in {psutil.STATUS_DEAD, psutil.STATUS_ZOMBIE}
+    except psutil.NoSuchProcess:
+      return False
+
+  def _killpid(self, pid, kill_timeout=None):
+    try:
+      proc = psutil.Process(pid)
+      proc.terminate()
+      time.sleep(kill_timeout or 1.0)
+      proc.kill()
+    except psutil.NoSuchProcess:
+      pass
+
+  def _setup_daemon(self, wpipe):
+    # This 2nd os.setsid() makes the daemon a process group, so with can kill the
+    # whole group, if required. We do this only on systems that supports it.
+    if os.name == 'posix':
+      os.setsid()
+
+    sys.stdout.flush()
+    sys.stderr.flush()
+
+    infd = os.open(os.devnull, os.O_RDONLY)
+    outfd = os.open(os.devnull, os.O_WRONLY | os.O_APPEND)
+    errfd = os.open(os.devnull, os.O_WRONLY | os.O_APPEND)
+
+    os.dup2(infd, sys.stdin.fileno())
+    os.dup2(outfd, sys.stdout.fileno())
+    os.dup2(errfd, sys.stderr.fileno())
+
+    pid = os.getpid()
+    self._writepid(pid)
+
+    # Register the signal handlers otherwise atexit callbacks will not get
+    # called in case a signal terminates the daemon process.
+    signal.signal(signal.SIGINT, _term_handler)
+    signal.signal(signal.SIGTERM, _term_handler)
+    atexit.register(functools.partial(self._delpid, pid=pid))
+
+    self._write_result(wpipe, pid=pid)
+
+  def stop(self, kill_timeout=None):
+    with self._lockfile():
+      pid = self._readpid()
+      if pid is not None:
+        self._killpid(pid, kill_timeout=kill_timeout)
+
+    self._delpid(pid=pid)
+
+    return pid is not None
+
+
+class DaemonPosix(DaemonBase):
+
+  def _daemonize(self):
+    rpipe, wpipe = os.pipe()
+    os.set_inheritable(wpipe, True)
+
+    pid = os.fork()
+    if pid > 0:
+      dres = self._read_result(rpipe)
+      if dres.pid < 0:
+        raise dres.exclass(dres.msg)
+
+      return dres.pid
+
+    try:
+      os.chdir('/')
+      os.setsid()
+      os.umask(0)
+
+      pid = os.fork()
+      if pid > 0:
+        sys.exit(0)
+
+      self._setup_daemon(wpipe)
+
+      return 0
+    except Exception as ex:
+      self._write_result(wpipe, exclass=ex.__class__, msg=f'Daemonize failed: {ex}')
+      sys.exit(1)
+
+  def start(self, target):
+    pid = self.getpid()
+    if pid is None:
+      if (pid := self._daemonize()) == 0:
+        target()
+        sys.exit(0)
+
+    return pid
+
+
+class DaemonCompat(DaemonBase):
+
+  def _write_result(self, wpipe, **kwargs):
+    dres = DaemonResult(**kwargs)
+    wpipe.send(pickle.dumps(dres))
+
+  def _read_result(self, rpipe):
+    return pickle.loads(rpipe.recv())
+
+  def _boostrap(self, target, wpipe):
+    try:
+      os.chdir('/')
+      if os.name == 'posix':
+        os.umask(0)
+
+      self._setup_daemon(wpipe)
+    except Exception as ex:
+      self._write_result(wpipe, exclass=ex.__class__, msg=f'Daemonize failed: {ex}')
+      sys.exit(1)
+
+    target()
+    sys.exit(0)
+
+  def _start_daemon(self, target):
+    mps = multiprocessing.get_context(method=os.getenv('DAEMON_MP_CONTEXT'))
+    rpipe, wpipe = mps.Pipe()
+    proc = mps.Process(target=self._boostrap, args=(target, wpipe))
+    proc.start()
+
+    dres = self._read_result(rpipe)
+    if dres.pid < 0:
+      raise dres.exclass(dres.msg)
+
+    assert dres.pid == proc.pid
+
+    # HACK!
+    # This removes the daemon process from the list of known child processes, preventing
+    # it to be forcibly killed at exit.
+    multiprocessing.process._children.discard(proc)
+
+    return dres.pid
+
+  def start(self, target):
+    pid = self.getpid()
+    if pid is None:
+      pid = self._start_daemon(target)
+
+    return pid
+
+
+try:
+  HAS_MP_CHILDREN = isinstance(multiprocessing.process._children, set)
+except:
+  HAS_MP_CHILDREN = False
+
+if HAS_MP_CHILDREN and os.getenv('DAEMON_MODE') != 'posix':
+  Daemon = DaemonCompat
+else:
+  Daemon = DaemonPosix
+

py_misc_utils/data_cache.py

@@ -0,0 +1,46 @@
+import hashlib
+import os
+import pickle
+import time
+
+from . import gfs
+
+
+_MAX_AGE = int(os.getenv('DATACACHE_MAX_AGE', 3600 * 24))
+
+class DataCache:
+
+  def __init__(self, data_id, max_age=_MAX_AGE):
+    fsid = hashlib.sha1(data_id.encode()).hexdigest()
+
+    cache_path = os.path.join(gfs.cache_dir(), 'data_cache')
+    gfs.makedirs(cache_path, exist_ok=True)
+
+    self._data_path = os.path.join(cache_path, fsid)
+    self._orig_data = None
+    try:
+      sres = gfs.stat(self._data_path)
+      if max_age is None or (sres.st_ctime + max_age) > time.time():
+        with gfs.open(self._data_path, mode='rb') as fd:
+          self._orig_data = pickle.load(fd)
+    except:
+      pass
+
+    self._data = self._orig_data
+
+  def data(self):
+    return self._data
+
+  def store(self, data):
+    self._data = data
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, *exc):
+    if self._data is not None and self._data is not self._orig_data:
+      with gfs.open(self._data_path, mode='wb') as fd:
+        pickle.dump(self._data, fd)
+
+    return False
+

py_misc_utils/date_utils.py

@@ -0,0 +1,90 @@
+import datetime
+import os
+import pytz
+import re
+
+import dateutil
+import dateutil.parser
+import numpy as np
+import pandas as pd
+
+from . import assert_checks as tas
+
+
+# Default timezone so code shows/works with such timezone.
+DEFAULT_TZ = pytz.timezone(os.getenv('DEFAULT_TZ', 'America/New_York'))
+
+
+def ny_market_timezone():
+  return pytz.timezone('America/New_York')
+
+
+def now(tz=None):
+  return datetime.datetime.now(tz=tz or DEFAULT_TZ)
+
+
+def from_timestamp(ts, tz=None):
+  return datetime.datetime.fromtimestamp(ts, tz=tz or DEFAULT_TZ)
+
+
+def make_datetime_from_epoch(s, tz=None):
+  ds = pd.to_datetime(s, unit='s', origin='unix', utc=True)
+
+  if tz is not None:
+    return ds.dt.tz_convert(tz) if isinstance(ds, pd.Series) else ds.tz_convert(tz)
+
+  return ds
+
+
+def parse_date(dstr, tz=None):
+  # Accept EPOCH values starting with @.
+  m = re.match(r'@((\d+)(\.\d*)?)$', dstr)
+  if m:
+    return from_timestamp(float(m.group(1)), tz=tz)
+
+  # ISO Format: 2011-11-17T00:05:23-04:00
+  dt = dateutil.parser.isoparse(dstr)
+  if dt.tzinfo is None and tz is not None:
+    dt = tz.localize(dt)
+
+  return dt
+
+
+def day_offset(dt):
+  ddt = dt.replace(hour=0, minute=0, second=0, microsecond=0)
+
+  return dt.timestamp() - ddt.timestamp(), ddt
+
+
+def np_datetime_to_epoch(dt, dtype=np.float64):
+  tas.check_fn(np.issubdtype, dt.dtype, np.datetime64)
+  u, c = np.datetime_data(dt.dtype)
+  if u == 's':
+    return dt.astype(dtype)
+
+  dt = dt.astype(np.float64)
+  if u == 'ns':
+    dt = dt / 1e9
+  elif u == 'us':
+    dt = dt / 1e6
+  elif u == 'ms':
+    dt = dt / 1e3
+  else:
+    alog.xraise(RuntimeError, f'Unknown NumPy datetime64 unit: {u}')
+
+  return dt if dtype == np.float64 else dt.astype(dtype)
+
+
+def align(dt, step, ceil=False):
+  secs = step.total_seconds() if isinstance(step, datetime.timedelta) else step
+  fp, ip = np.modf(dt.timestamp() / secs)
+
+  if ceil and not np.isclose(fp, 0):
+    ip += 1
+
+  return from_timestamp(ip * secs, tz=dt.tzinfo)
+
+
+def at_time(dt, hour=0, minute=0, second=0, microsecond=0):
+  return dt.replace(hour=hour, minute=minute, second=second, microsecond=microsecond)
+

py_misc_utils/debug.py

@@ -0,0 +1,24 @@
+import pdb
+import signal
+import sys
+
+
+def _debug(signum, frame):
+  signame = signal.strsignal(signum)
+  sys.stderr.write(f'** {signame} received, entering debugger\n' \
+                   f'** Type "c" to continue or "q" to stop the process\n' \
+                   f'** Or {signame} again to quit (and possibly dump core)\n')
+  sys.stderr.flush()
+
+  signal.signal(signum, signal.SIG_DFL)
+  try:
+    pdb.set_trace()
+  finally:
+    signal.signal(signum, _debug)
+
+
+def install_pdb_hook(signum):
+  signum = getattr(signal, signum) if isinstance(signum, str) else signum
+
+  signal.signal(signum, _debug)
+

py_misc_utils/dyn_modules.py

@@ -0,0 +1,50 @@
+import collections
+import importlib
+import os
+
+from . import alog
+from . import assert_checks as tas
+from . import gfs
+from . import module_utils as mu
+
+
+class DynLoader:
+
+  def __init__(self, modname=None, path=None, postfix=''):
+    if modname is not None:
+      tas.check_is_none(path, msg=f'Cannot specify path="{path}" when specified modname="{modname}"')
+
+      parent_mod = importlib.import_module(modname)
+      mpath = os.path.dirname(parent_mod.__file__)
+    else:
+      tas.check_is_not_none(path, msg=f'Path must be specified if "modname" is missing')
+      mpath = path
+
+    module_names, matcher = [], gfs.RegexMatcher(r'(.*)' + postfix + r'\.py$')
+    for fname in gfs.enumerate_files(mpath, matcher=matcher):
+      module_names.append((matcher.match.group(1), os.path.join(mpath, fname)))
+
+    self._modules = collections.OrderedDict()
+    for (imod_name, imod_path) in sorted(module_names):
+      if modname is not None:
+        imod = importlib.import_module(f'{modname}.{imod_name}{postfix}')
+      else:
+        imod = mu.load_module(imod_path, modname=f'{imod_name}{postfix}')
+      mname = getattr(imod, 'MODULE_NAME', imod_name)
+      self._modules[mname] = imod
+
+  def module_names(self):
+    return tuple(self._modules.keys())
+
+  def modules(self):
+    return tuple(self._modules.values())
+
+  def get(self, name):
+    return self._modules.get(name)
+
+  def __getitem__(self, name):
+    return self._modules[name]
+
+  def __len__(self):
+    return len(self._modules)
+

py_misc_utils/dynamod.py

@@ -0,0 +1,103 @@
+import hashlib
+import importlib
+import os
+import shutil
+import sys
+import threading
+
+from . import global_namespace as gns
+
+
+_MODROOT = 'pym'
+
+def _create_root():
+  from . import tempdir as tmpd
+
+  path = os.path.join(tmpd.get_temp_root(), _MODROOT)
+  os.mkdir(path)
+
+  return path
+
+
+_MODNAME = '_dynamod'
+
+def _create_mod_folder():
+  path = _create_root()
+  dpath = os.path.join(path, _MODNAME)
+  os.mkdir(dpath)
+
+  with open(os.path.join(dpath, '__init__.py'), mode='w') as fd:
+    pass
+
+  sys.path.append(path)
+
+  return dpath
+
+
+def _clone_mod_folder(src_path):
+  # This copies the source folder into a new temporary one for the new process,
+  # which will be in turn deleted once this exists.
+  # We cannot point directly to the source folder since it will be removed once
+  # the parent exists.
+  path = _create_root()
+  shutil.copytree(src_path, os.path.join(path, _MODNAME))
+
+  src_root = os.path.dirname(src_path)
+  if src_root in sys.path:
+    sys.path.remove(src_root)
+  sys.path.append(path)
+
+  return path
+
+
+_LOCK = threading.RLock()
+_MOD_FOLDER = gns.Var('dynamod.MOD_FOLDER',
+                      child_fn=_clone_mod_folder,
+                      defval=_create_mod_folder)
+
+_HASHNAME_LEN = int(os.getenv('DYNAMOD_HASHNAME_LEN', 12))
+
+def make_code_name(code):
+  chash = hashlib.sha1(code.encode()).hexdigest()[: _HASHNAME_LEN]
+
+  return f'_hashed._{chash}'
+
+
+def create_module(name, code, overwrite=None):
+  path = gns.get(_MOD_FOLDER)
+  mpath = os.path.join(path, *name.split('.')) + '.py'
+
+  with _LOCK:
+    reload = False
+    if os.path.exists(mpath):
+      if overwrite in (None, False):
+        raise RuntimeError(f'Dynamic module "{name}" already exists: {mpath}')
+      else:
+        # Note that there exist an issue with Python import subsystem:
+        #
+        #   https://bugs.python.org/issue31772
+        #
+        # Such issue causes a module to not be reloaded if the size of the source file
+        # has not changed.
+        # So BIG HACK here to add an headline comment if that's the case!
+        reload = True
+        if os.stat(mpath).st_size == len(code):
+          code = f'# Note: Added due to https://bugs.python.org/issue31772\n\n{code}'
+
+    os.makedirs(os.path.dirname(mpath), exist_ok=True)
+    with open(mpath, mode='w') as f:
+      f.write(code)
+
+    module = get_module(name)
+
+    return importlib.reload(module) if reload else module
+
+
+def module_name(name):
+  return f'{_MODNAME}.{name}'
+
+
+def get_module(name):
+  with _LOCK:
+    return importlib.import_module(module_name(name))
+

py_misc_utils/env_config.py

@@ -0,0 +1,35 @@
+import argparse
+
+from . import utils as ut
+
+
+class EnvConfig:
+
+  def __init__(self):
+    parser = argparse.ArgumentParser()
+    state = dict()
+    for name in dir(self):
+      if not name.startswith('_'):
+        value = getattr(self, name)
+        # Do not try to override functions (even though there really should not
+        # be in an EnvConfig derived object).
+        if not callable(value):
+          state[name] = ut.getenv(name, dtype=type(value))
+          parser.add_argument(f'--{name}', type=type(value))
+
+    args, _ = parser.parse_known_args()
+    for name, value in state.items():
+      avalue = getattr(args, name, None)
+      avalue = value if avalue is None else avalue
+      if avalue is not None:
+        setattr(self, name, avalue)
+
+  def __repr__(self):
+    cvars = dict()
+    for name in dir(self):
+      if not name.startswith('_'):
+        cvars[name] = getattr(self, name)
+
+    return ut.stri(cvars)
+
+