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

py_misc_utils/gfs.py

@@ -0,0 +1,371 @@
+import collections
+import contextlib
+import importlib
+import os
+import pkgutil
+import re
+import shutil
+import sys
+import urllib.parse as uparse
+
+from . import alog
+from . import assert_checks as tas
+from . import cached_file as chf
+from . import context_managers as cm
+from . import fs_utils as fsu
+from . import mirror_from as mrf
+from . import run_once as ro
+
+
+class TempFile:
+
+  def __init__(self, nsdir=None, nspath=None, **kwargs):
+    nsdir = nsdir if nsdir is None or is_local_path(nsdir) else None
+    nspath = nspath if nspath is None or is_local_path(nspath) else None
+
+    self._fs, self._path = resolve_fs(fsu.temp_path(nspath=nspath, nsdir=nsdir), **kwargs)
+    self._kwargs = kwargs
+    self._fd, self._delete = None, False
+
+  def open(self):
+    self._fd = self._fs.open(self._path, **self._kwargs)
+    self._delete = True
+    mrf.mirror_all(self._fd, self, name='fd')
+
+    return self
+
+  def _close_fd(self):
+    if self._fd is not None:
+      self._fd.close()
+      mrf.unmirror(self, name='fd')
+      self._fd = None
+
+  def close(self):
+    self._close_fd()
+    if self._delete:
+      self._fs.remove(self._path)
+      self._delete = False
+
+  def replace(self, path):
+    self._close_fd()
+    replace(self._path, path, src_fs=self._fs)
+    self._delete = False
+
+  def __enter__(self):
+    self.open()
+
+    return self
+
+  def __exit__(self, *exc):
+    self.close()
+
+    return False
+
+
+_STD_FILES = {
+  'STDIN': sys.stdin,
+  'STDOUT': sys.stdout,
+  'STDERR': sys.stderr,
+}
+
+def std_open(path, **kwargs):
+  if isinstance(path, str) and (sfd := _STD_FILES.get(path)) is not None:
+    return contextlib.nullcontext(sfd)
+
+  return open(path, **kwargs)
+
+
+def open(source, **kwargs):
+  if (path := path_of(source)) is not None:
+    fs, fpath = resolve_fs(path, **kwargs)
+
+    return fs.open(fpath, **kwargs)
+
+  return contextlib.nullcontext(source)
+
+
+def open_local(path, **kwargs):
+  return open(path, **kwargs)
+
+
+def maybe_open(path, **kwargs):
+  try:
+    return open(path, **kwargs)
+  except:
+    pass
+
+
+def as_local(path, **kwargs):
+  fs, fpath = resolve_fs(path, **kwargs)
+
+  return fs.as_local(fpath, **kwargs)
+
+
+def path_of(path):
+  return os.fspath(path) if isinstance(path, (str, os.PathLike)) else None
+
+
+PathSplit = collections.namedtuple('PathSplit', 'base, ext, purl')
+
+def splitext(path):
+  purl = uparse.urlparse(path)
+  base, ext = os.path.splitext(purl.path)
+
+  return PathSplit(base=base, ext=ext[1:], purl=purl)
+
+
+def is_file(path):
+  fs, fpath = resolve_fs(path)
+
+  return fs.isfile(fpath)
+
+
+def is_dir(path):
+  fs, fpath = resolve_fs(path)
+
+  return fs.isdir(fpath)
+
+
+def exists(path):
+  fs, fpath = resolve_fs(path)
+
+  return fs.exists(fpath)
+
+
+def is_same_fs(*args):
+  specs = []
+  for fspath in args:
+    purl = uparse.urlparse(fspath.path)
+    if purl.scheme:
+      specs.append((purl.scheme, purl.netloc))
+    else:
+      specs.append((_DEFAULT_LOCAL_PROTO, fsu.localfs_mount(purl.path)))
+
+  return all(specs[0] == s for s in specs[1:])
+
+
+_DEFAULT_LOCAL_PROTO = 'file'
+
+def is_local_proto(proto):
+  return proto == _DEFAULT_LOCAL_PROTO
+
+
+def is_local_fs(fs):
+  return is_local_proto(fs.ID)
+
+
+def is_local_path(path):
+  return is_local_proto(get_proto(path))
+
+
+def is_path(path):
+  # It is a path is it contains a proto, or starts with '/', './' or '../'.
+  return has_proto(path) or re.match(r'/|\.\.?/', path) is not None
+
+
+def has_proto(path):
+  return re.match(r'\w+://', path) is not None
+
+
+def get_proto(path):
+  m = re.match(r'(\w+)://', path)
+
+  return m.group(1).lower() if m else _DEFAULT_LOCAL_PROTO
+
+
+FsPath = collections.namedtuple('FsPath', 'fs, path')
+
+def resolve_paths(*paths):
+  resolved = []
+  for path_arg in paths:
+    if isinstance(path_arg, (list, tuple)):
+      fs, path = path_arg
+    else:
+      fs, path = None, path_arg
+    if fs is None:
+      fs, path = resolve_fs(path)
+
+    resolved.append(FsPath(fs, path))
+
+  return tuple(resolved)
+
+
+def copy(src_path, dest_path, src_fs=None, dest_fs=None):
+  src, dest = resolve_paths((src_fs, src_path), (dest_fs, dest_path))
+
+  src.fs.copyfile(src.path, dest.fs, dest.path)
+
+
+def replace(src_path, dest_path, src_fs=None, dest_fs=None):
+  src, dest = resolve_paths((src_fs, src_path), (dest_fs, dest_path))
+
+  if is_same_fs(src, dest):
+    dest.fs.replace(src.path, dest.path)
+  else:
+    copy(src.path, dest.path, src_fs=src.fs, dest_fs=dest.fs)
+    src.fs.remove(src.path)
+
+
+def remove(path):
+  fs, fpath = resolve_fs(path)
+  fs.remove(fpath)
+
+
+def mkdir(path, **kwargs):
+  fs, fpath = resolve_fs(path)
+  fs.mkdir(fpath, **kwargs)
+
+
+def makedirs(path, **kwargs):
+  fs, fpath = resolve_fs(path)
+  fs.makedirs(fpath, **kwargs)
+
+
+def rmdir(path):
+  fs, fpath = resolve_fs(path)
+  fs.rmdir(fpath)
+
+
+def rmtree(path, **kwargs):
+  fs, fpath = resolve_fs(path)
+  fs.rmtree(fpath, **kwargs)
+
+
+def stat(path):
+  fs, fpath = resolve_fs(path)
+
+  return fs.stat(fpath)
+
+
+def link(src_path, dest_path):
+  src, dest = resolve_paths(src_path, dest_path)
+
+  tas.check(is_same_fs(src, dest),
+            msg=f'Unable to link across file systems: {src_path} -> {dest_path}')
+
+  src.fs.link(src.path, dest.path)
+
+
+def symlink(src_path, dest_path):
+  src, dest = resolve_paths(src_path, dest_path)
+
+  tas.check(is_same_fs(src, dest),
+            msg=f'Unable to symlink across file systems: {src_path} -> {dest_path}')
+
+  src.fs.symlink(src.path, dest.path)
+
+
+class RegexMatcher:
+
+  def __init__(self, rex):
+    self._rex = re.compile(rex)
+    self.match = None
+
+  def __call__(self, value):
+    self.match = re.match(self._rex, value)
+
+    return self.match is not None
+
+
+def enumerate_files(path, matcher=None, return_stats=False):
+  fs, fpath = resolve_fs(path)
+
+  for de in fs.list(fpath):
+    if matcher is None or matcher(de.name):
+      if return_stats:
+        yield de.name, de
+      else:
+        yield de.name
+
+
+def normpath(path):
+  _, fpath = resolve_fs(path)
+
+  return fpath
+
+
+_CACHE_DIR = fsu.normpath(os.getenv('CACHE_DIR',
+                                    os.path.join(fsu.home(), '.cache')))
+
+def cache_dir():
+  return _CACHE_DIR
+
+
+def set_cache_dir(path):
+  global _CACHE_DIR
+
+  _CACHE_DIR = fsu.normpath(path)
+
+
+def find_mount(path):
+  fs, fpath = resolve_fs(path)
+
+  return fsu.localfs_mount(fpath) if is_local_fs(fs) else None
+
+
+_FS_REGISTRY = dict()
+
+def register_fs(cls):
+  for fsid in cls.IDS:
+    alog.debug(f'Registering file system: {fsid}')
+    _FS_REGISTRY[fsid] = cls
+
+
+def try_register(importer, modname, parent=None):
+  try:
+    if parent is None:
+      spec = importer.find_spec(modname)
+      module = importlib.util.module_from_spec(spec)
+      spec.loader.exec_module(module)
+    else:
+      module = importlib.import_module(f'{parent}.{modname}')
+
+    file_systems = getattr(module, 'FILE_SYSTEMS', ())
+    for cls in file_systems:
+      register_fs(cls)
+
+    return module
+  except ImportError as ex:
+    alog.verbose(f'Unable to import file system module "{modname}": {ex}')
+
+
+def register_fs_from_path(path, parent=None):
+  for importer, modname, _ in pkgutil.iter_modules(path=path):
+    if modname.endswith('_fs'):
+      try_register(importer, modname, parent=parent)
+
+
+@ro.run_once
+def register_modules():
+  import py_misc_utils.fs as pyfs
+
+  register_fs_from_path(pyfs.__path__, parent='py_misc_utils.fs')
+
+  gfs_path = os.getenv('GFS_PATH')
+  if gfs_path:
+    for path in gfs_path.split(':'):
+      register_fs_from_path(path)
+
+
+def get_proto_fs(proto, **kwargs):
+  register_modules()
+
+  cls = _FS_REGISTRY.get(proto)
+  tas.check_is_not_none(cls, msg=f'Protocol "{proto}" not registered')
+
+  return cls(**kwargs)
+
+
+def resolve_fs(path, **kwargs):
+  proto = get_proto(path)
+
+  cachedir = chf.get_cache_dir(kwargs.pop('cache_dir', cache_dir()))
+
+  cache_iface = kwargs.pop('cache_iface', None)
+  if cache_iface is None:
+    cache_iface = chf.CacheInterface(cachedir)
+
+  fs = get_proto_fs(proto, cache_iface=cache_iface, cache_dir=cachedir, **kwargs)
+
+  return fs, fs.norm_url(path)
+

py_misc_utils/git_repo.py

@@ -0,0 +1,77 @@
+import os
+import shutil
+import subprocess
+
+from . import alog
+from . import assert_checks as tas
+
+
+class GitRepo:
+
+  def __init__(self, path):
+    self.path = path
+
+  def _git(self, *cmd):
+    git_cmd = ['git', '-C', self.path] + list(cmd)
+    alog.debug(f'Running GIT: {git_cmd}')
+
+    return git_cmd
+
+  def _run(self, *cmd):
+    rcmd = []
+    for arg in cmd:
+      if isinstance(arg, (list, tuple)):
+        rcmd.extend(arg)
+      else:
+        rcmd.appen(arg)
+
+    subprocess.run(rcmd, capture_output=True, check=True)
+
+  def _cmd(self, *cmd):
+    self._run(self._git(*cmd))
+
+  def _outcmd(self, *cmd, strip=False):
+    output = subprocess.check_output(self._git(*cmd))
+    if isinstance(output, bytes):
+      output = output.decode()
+
+    return output.strip() if strip else output
+
+  def repo(self):
+    return self._outcmd('config', '--get', 'remote.origin.url', strip=True)
+
+  def clone(self, repo, force=False, shallow=False):
+    do_clone = True
+    if os.path.isdir(self.path):
+      tas.check_eq(repo, self.repo(), msg=f'Repo mismatch!')
+      if force or shallow != self.is_shallow():
+        alog.info(f'Purging old GIT folder: {self.path}')
+        shutil.rmtree(self.path)
+      else:
+        self.pull()
+        do_clone = False
+
+    if do_clone:
+      parent_path = os.path.dirname(self.path)
+      os.makedirs(parent_path, exist_ok=True)
+      git_cmd = ['git', '-C', parent_path, 'clone', '-q']
+      if shallow:
+        git_cmd += ['--depth', '1', repo, os.path.basename(self.path)]
+      else:
+        git_cmd += [repo, os.path.basename(self.path)]
+
+      alog.debug(f'Running GIT: {git_cmd}')
+      self._run(git_cmd)
+
+  def current_commit(self):
+    return self._outcmd('rev-parse', 'HEAD', strip=True)
+
+  def is_shallow(self):
+    return self._outcmd('rev-parse', '--is-shallow-repository', strip=True) == 'true'
+
+  def pull(self):
+    self._cmd('pull', '-q')
+
+  def checkout(self, commit):
+    self._cmd('checkout', '-q', commit)
+

py_misc_utils/global_namespace.py

@@ -0,0 +1,110 @@
+# When using multiprocessing, there are two distinct behaviours if fork start method
+# is used, WRT spawn/forkserver.  In the latter case the global context
+# accumulated by the running process is not pickled-through the child, so the new
+# process start with wiped out global namespace.
+# Using this API, together with the multiprocessing.create_process(), it is possible to
+# have global data transfered to the child.
+# All data stored in the global namespace must be pickle-able, unless fork_init is
+# set to True.
+# If the fork_init attribute is True, it means the variables data must be cleared
+# within the child process, and not carried over (COW-ed) like it would happen when
+# using the fork(2) system call.
+# NOTE: This is a low level module which should have no explicit local dependencies.
+
+import collections
+import inspect
+import multiprocessing
+import os
+import threading
+
+
+# The parent_fn function is called (if present) before the creation of a new process,
+# within the parent, with the current value of the variable, and is supposed to be
+# returning the "state" of such variable. The state must be pickle-able.
+# The child_fn function is called (if present) after the creation of a new process,
+# within the child, to restore a variable from its state (the value of the new variable
+# should be returned).
+Var = collections.namedtuple(
+  'Var',
+  'name, parent_fn, child_fn, data, fork_init, defval',
+  defaults=(None, None, None, False, None))
+
+_NS = dict()
+_LOCK = threading.RLock()
+
+
+def _child_fork():
+  global _NS, _LOCK
+
+  cns = dict()
+  for var in _NS.values():
+    if not var.fork_init:
+      cns[var.name] = var
+
+  _NS = cns
+  _LOCK = threading.RLock()
+
+
+if os.name == 'posix':
+  os.register_at_fork(after_in_child=_child_fork)
+
+
+def parent_switch(method):
+  assert method in multiprocessing.get_all_start_methods(), method
+
+  pns = dict()
+  with _LOCK:
+    for var in _NS.values():
+      # Variables with fork_init=True are the ones that are supposed to be
+      # initialized in every process, and as such do not have to be carried over
+      # from the parent context. Also, fork_init=True variables might contain data
+      # which is not pickle-able, and carrying them over will fail.
+      if not var.fork_init:
+        if var.parent_fn is not None:
+          data = var.parent_fn(var.data)
+          if data is not var.data:
+            var = None if data is None else var._replace(data=data)
+
+        if var is not None:
+          pns[var.name] = var
+
+  return pns
+
+
+def child_switch(method, ns):
+  global _NS
+
+  assert method in multiprocessing.get_all_start_methods(), method
+
+  cns = dict()
+  for var in ns.values():
+    if var.child_fn is not None:
+      data = var.child_fn(var.data)
+      if data is not var.data:
+        var = None if data is None else var._replace(data=data)
+
+    if var is not None:
+      cns[var.name] = var
+
+  _NS = cns
+
+
+def get(var, force=True):
+  with _LOCK:
+    value = _NS.get(var.name)
+    if value is None and force:
+      data = var.defval() if inspect.isfunction(var.defval) else var.defval
+      value = var._replace(data=data)
+      _NS[value.name] = value
+
+  return value.data if value is not None else None
+
+
+def set(var, data):
+  with _LOCK:
+    prev_value = _NS.get(var.name)
+    value = var._replace(data=data)
+    _NS[value.name] = value
+
+  return prev_value.data if prev_value is not None else None
+

py_misc_utils/http_async_fetcher.py

@@ -0,0 +1,139 @@
+import functools
+import os
+
+import httpx
+
+from . import async_manager as asym
+from . import assert_checks as tas
+from . import core_utils as cu
+from . import file_overwrite as fow
+from . import fin_wrap as fw
+from . import gfs
+from . import tempdir as tmpd
+from . import utils as ut
+from . import work_results as wres
+
+
+async def http_fetch_url(url, context=None, path=None, http_args=None):
+  wpath = wres.work_path(path, url)
+  try:
+    client = await context.get('httpx.AsyncClient', httpx.AsyncClient)
+
+    resp = await client.get(url, **http_args)
+    resp.raise_for_status()
+
+    with wres.write_result(wpath) as fd:
+      fd.write(resp.content)
+  except Exception as ex:
+    wres.write_error(wpath, ex, workid=url)
+  finally:
+    return wpath
+
+
+class HttpAsyncFetcher:
+
+  def __init__(self,
+               path=None,
+               num_workers=None,
+               http_args=None,
+               mpctx=None):
+    self._ctor_path = path
+    self._path = None
+    self._num_workers = num_workers
+    self._http_args = ut.dict_setmissing(
+      http_args or dict(),
+      timeout=ut.getenv('FETCHER_TIMEO', dtype=float, defval=10.0),
+    )
+    self._mpctx = mpctx
+    self._async_manager = None
+    self._pending = set()
+
+  @classmethod
+  def _cleaner(cls, self):
+    self._async_manager.close()
+    if self._path != self._ctor_path:
+      gfs.rmtree(self._path, ignore_errors=True)
+
+  def start(self):
+    if self._ctor_path is None:
+      self._path = tmpd.fastfs_dir()
+    else:
+      self._path = self._ctor_path
+
+    async_manager = asym.AsyncManager(**cu.denone(num_workers=self._num_workers,
+                                                  mpctx=self._mpctx))
+
+    finfn = functools.partial(self._cleaner,
+                              cu.object_context(self, _async_manager=async_manager))
+    fw.fin_wrap(self, '_async_manager', async_manager, finfn=finfn)
+
+  def shutdown(self):
+    async_manager = self._async_manager
+    if async_manager is not None:
+      fw.fin_wrap(self, '_async_manager', None, cleanup=True)
+      self._path = None
+      self._pending = set()
+
+  def enqueue(self, *urls):
+    wmap = dict()
+    for url in urls:
+      if url:
+        work_ctor = functools.partial(http_fetch_url, url,
+                                      path=self._path,
+                                      http_args=self._http_args)
+        self._async_manager.enqueue_work(url, work_ctor)
+        self._pending.add(url)
+        wmap[url] = wres.work_hash(url)
+
+    return wmap
+
+  def wait(self, url):
+    wpath = wres.work_path(self._path, url)
+    if not os.path.isfile(wpath):
+      tas.check(url in self._pending, msg=f'URL already retired: {url}')
+
+      while self._pending:
+        (rurl, result) = self._async_manager.fetch_result()
+
+        self._pending.discard(rurl)
+        wres.raise_if_error(result)
+        if rurl == url:
+          break
+
+    try:
+      return wres.get_work(wpath)
+    finally:
+      os.remove(wpath)
+
+  def iter_results(self, max_results=None, block=True, timeout=None):
+    count = 0
+    while self._pending:
+      if (fetchres := self._async_manager.fetch_result(block=block,
+                                                       timeout=timeout)) is None:
+        break
+
+      rurl, result = fetchres
+
+      self._pending.discard(rurl)
+      wpath = wres.work_path(self._path, rurl)
+
+      wdata = wres.load_work(wpath)
+
+      os.remove(wpath)
+
+      yield rurl, wdata
+
+      count += 1
+      if max_results is not None and count >= max_results:
+        break
+
+  def __enter__(self):
+    self.start()
+
+    return self
+
+  def __exit__(self, *exc):
+    self.shutdown()
+
+    return False
+