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

py_misc_utils/fs_utils.py

@@ -0,0 +1,207 @@
+import contextlib
+import functools
+import io
+import os
+import pathlib
+import shutil
+import tempfile
+
+from . import alog
+from . import assert_checks as tas
+from . import osfd
+from . import rnd_utils as rngu
+
+
+def home():
+  return pathlib.Path.home()
+
+
+def maybe_remove(path):
+  if os.path.exists(path):
+    os.remove(path)
+
+
+def is_binary(fd):
+  if (mode := getattr(fd, 'mode', None)) is not None:
+    return 'b' in mode
+
+  return not isinstance(fd, io.TextIOBase)
+
+
+def link_or_copy(src_path, dest_path):
+  try:
+    os.link(src_path, dest_path)
+
+    return dest_path
+  except OSError as ex:
+    alog.debug(f'Harklink failed from "{src_path}" to "{dest_path}", trying symlink. ' \
+               f'Error was: {ex}')
+
+  try:
+    os.symlink(src_path, dest_path)
+
+    return dest_path
+  except OSError:
+    alog.debug(f'Symlink failed from "{src_path}" to "{dest_path}", going to copy. ' \
+               f'Error was: {ex}')
+
+  shutil.copyfile(src_path, dest_path)
+  shutil.copystat(src_path, dest_path)
+
+  return dest_path
+
+
+def is_newer_file(path, other):
+  return os.stat(path).st_mtime > os.stat(other).st_mtime
+
+
+def os_opener(*args, **kwargs):
+  return functools.partial(os.open, *args, **kwargs)
+
+
+def safe_rmtree(path, **kwargs):
+  tpath = temp_path(nspath=path)
+  try:
+    os.rename(path, tpath)
+  except FileNotFoundError:
+    if not kwargs.get('ignore_errors', False):
+      raise
+  else:
+    shutil.rmtree(tpath, **kwargs)
+
+
+def readall(fd):
+  if isinstance(fd, str):
+    with osfd.OsFd(fd, os.O_RDONLY) as ifd:
+      return readall(ifd)
+  else:
+    sres = os.stat(fd)
+    os.lseek(fd, 0, os.SEEK_SET)
+
+    return os.read(fd, sres.st_size)
+
+
+def stat(path):
+  try:
+    return os.stat(path)
+  except:
+    pass
+
+
+def path_split(path):
+  path_parts = []
+  while True:
+    parts = os.path.split(path)
+    if parts[0] == path:
+      path_parts.append(parts[0])
+      break
+    elif parts[1] == path:
+      path_parts.append(parts[1])
+      break
+    else:
+      path = parts[0]
+      path_parts.append(parts[1])
+
+  path_parts.reverse()
+
+  return path_parts
+
+
+def drop_ext(path, exts):
+  xpath, ext = os.path.splitext(path)
+
+  if isinstance(exts, str):
+    return xpath if ext == exts else path
+
+  return xpath if ext in exts else path
+
+
+def normpath(path):
+  path = os.path.expanduser(path)
+  path = os.path.expandvars(path)
+  path = os.path.abspath(path)
+
+  return os.path.normpath(path)
+
+
+def localfs_mount(path):
+  while True:
+    parent_path = os.path.dirname(path)
+    if path == parent_path or os.path.ismount(path):
+      return path
+    path = parent_path
+
+
+def enum_chunks(stream, chunk_size=16 * 1024**2):
+  while True:
+    data = stream.read(chunk_size)
+    if data:
+      yield data
+    if chunk_size > len(data):
+      break
+
+
+def du(path, follow_symlinks=None, visited=None):
+  visited = set() if visited is None else visited
+  follow_symlinks = follow_symlinks not in (None, False)
+
+  size, dirs = 0, []
+  if path not in visited:
+    visited.add(path)
+    with os.scandir(path) as sdit:
+      for de in sdit:
+        if de.is_file(follow_symlinks=follow_symlinks):
+          sres = de.stat()
+          size += sres.st_size
+        elif de.is_dir(follow_symlinks=follow_symlinks):
+          dirs.append(de.name)
+
+    for dname in dirs:
+      size += du(os.path.join(path, dname),
+                 follow_symlinks=follow_symlinks,
+                 visited=visited)
+
+  return size
+
+
+def find_path(name, paths, checkfn=os.path.exists):
+  for path in paths:
+    cpath = os.path.join(path, name)
+    if checkfn(cpath):
+      return cpath
+
+
+_TMPFN_RNDSIZE = int(os.getenv('TMPFN_RNDSIZE', 10))
+
+def temp_path(nspath=None, nsdir=None, rndsize=_TMPFN_RNDSIZE):
+  if nspath is not None:
+    bpath, ext = os.path.splitext(nspath)
+
+    return f'{bpath}.{rngu.rand_string(rndsize)}{ext}'
+
+  nsdir = tempfile.gettempdir() if nsdir is None else nsdir
+
+  return os.path.join(nsdir, f'{rngu.rand_string(rndsize)}.tmp')
+
+
+# This does FileOverwrite() task (although locally limited) but here we do not
+# pull that dependency to allow inlcusion in this module (which allows none).
+@contextlib.contextmanager
+def atomic_write(path, mode='wb', create_parents=False):
+  tpath = temp_path(nspath=path)
+
+  if create_parents:
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+  fd = open(tpath, mode=mode)
+  try:
+    yield fd
+    fd.close()
+    fd = None
+  finally:
+    if fd is not None:
+      fd.close()
+      os.remove(tpath)
+    else:
+      os.replace(tpath, path)
+

py_misc_utils/gcs_fs.py

@@ -0,0 +1,169 @@
+import collections
+import os
+import stat as st
+
+import google.cloud.storage as gcs
+
+from . import alog
+from . import fs_base as fsb
+
+
+# https://cloud.google.com/appengine/docs/legacy/standard/python/googlecloudstorageclient/read-write-to-cloud-storage
+# https://cloud.google.com/python/docs/reference/storage/1.44.0/client
+# https://cloud.google.com/python/docs/reference/storage/1.44.0/blobs#google.cloud.storage.blob.Blob
+
+class GcsFs:
+
+  def __init__(self, bucket):
+    self.bucket = bucket
+    self._client = gcs.Client()
+
+  def _blob_stat(self, blob, base_path=None):
+    name = blob.name
+    if base_path is not None:
+      if not name.startswith(base_path):
+        return
+      name = name[len(base_path):]
+      spos = name.find('/')
+      if spos > 0:
+        name = name[: spos]
+        mode = st.S_IFDIR
+        size, etag = 0, None
+      else:
+        mode = st.S_IFREG
+        size, etag = blob.size, blob.etag
+
+      path = base_path + name
+    else:
+      spos = name.rfind('/')
+      if spos >= 0:
+        name = name[spos + 1:]
+      mode = st.S_IFREG
+      path = blob.name
+      size, etag = blob.size, blob.etag
+
+    return fsb.DirEntry(name=name,
+                        path=path,
+                        etag=etag,
+                        st_mode=mode,
+                        st_size=size,
+                        st_ctime=blob.time_created.timestamp(),
+                        st_mtime=blob.updated.timestamp())
+
+  def _norm_path(self, path):
+    if path:
+      if path == '/':
+        path = ''
+      else:
+        path = path + '/' if not path.endswith('/') else path
+
+    return path
+
+  def listdir(self, path):
+    npath = self._norm_path(path)
+
+    dentries = dict()
+    for blob in self._client.list_blobs(self.bucket, prefix=npath):
+      if (dentry := self._blob_stat(blob, base_path=npath)) is not None:
+        xdentry = dentries.get(dentry.name)
+        if xdentry is not None:
+          dentry = dentry._replace(st_ctime=min(dentry.st_ctime, xdentry.st_ctime),
+                                   st_mtime=max(dentry.st_mtime, xdentry.st_mtime))
+
+        dentries[dentry.name] = dentry
+
+    sorted_dentries = sorted(dentries.items(), key=lambda x: (x[1].st_mode, x[0]))
+    for name, dentry in sorted_dentries:
+      yield dentry
+
+  def open(self, path, mode='rb'):
+    bucket = self._client.bucket(self.bucket)
+    blob = bucket.blob(path)
+
+    return blob.open(mode)
+
+  def upload(self, path, source):
+    bucket = self._client.bucket(self.bucket)
+    blob = bucket.blob(path)
+
+    with blob.open('wb') as fd:
+      for data in source:
+        fd.write(data)
+
+  def download(self, path, chunk_size=32 * 1024**2):
+    bucket = self._client.bucket(self.bucket)
+    blob = bucket.blob(path)
+
+    with blob.open('rb') as fd:
+      while True:
+        data = fd.read(chunk_size)
+        if data:
+          yield data
+        if chunk_size > len(data):
+          break
+
+  def pread(self, path, offset, size):
+    bucket = self._client.bucket(self.bucket)
+    blob = bucket.blob(path)
+
+    return blob.download_as_bytes(start=offset, end=offset + size - 1, raw_download=True)
+
+  def exists(self, path):
+    bucket = self._client.bucket(self.bucket)
+    blob = bucket.blob(path)
+
+    return blob.exists()
+
+  def stat(self, path):
+    bucket = self._client.bucket(self.bucket)
+    blob = bucket.get_blob(path)
+    if blob is not None:
+      return self._blob_stat(blob)
+
+    ctime = mtime = None
+    for de in self.listdir(path):
+      if ctime is None or de.st_ctime < ctime:
+        ctime = de.st_ctime
+      if mtime is None or de.st_mtime > mtime:
+        mtime = de.st_mtime
+
+    if ctime is not None and mtime is not None:
+      bpath = path[: -1] if path.endswith('/') else path
+      name = os.path.basename(bpath)
+
+      return fsb.DirEntry(name=name,
+                          path=bpath,
+                          st_mode=st.S_IFDIR,
+                          st_size=0,
+                          st_ctime=ctime,
+                          st_mtime=mtime)
+
+  def remove(self, path):
+    bucket = self._client.bucket(self.bucket)
+    blob = bucket.blob(path)
+    blob.delete()
+
+  def rename(self, src_path, dest_path):
+    bucket = self._client.bucket(self.bucket)
+    src_blob = bucket.blob(src_path)
+
+    bucket.copy_blob(src_blob, bucket, dest_path)
+    bucket.delete_blob(src_path)
+
+  def rmtree(self, path, ignore_errors=None):
+    npath = self._norm_path(path)
+
+    for blob in self._client.list_blobs(self.bucket, prefix=npath):
+      try:
+        blob.delete()
+      except Exception as ex:
+        alog.debug(f'Failed to remove "{blob.name}" from "{self.bucket}": {ex}')
+        if ignore_errors in (None, False):
+          raise
+
+  def copy(self, src_path, dest_path):
+    bucket = self._client.bucket(self.bucket)
+    src_blob = bucket.blob(src_path)
+
+    bucket.copy_blob(src_blob, bucket, dest_path)
+

py_misc_utils/gen_indices.py

@@ -0,0 +1,54 @@
+import itertools
+
+import numpy as np
+
+
+def generate(shape, idx):
+  if isinstance(idx, int):
+    assert len(shape) == 1, f'{shape}'
+
+    gens = [(idx,)]
+  elif isinstance(idx, slice):
+    assert len(shape) == 1, f'{shape}'
+
+    gens = [range(*idx.indices(shape[0]))]
+  elif isinstance(idx, (list, np.ndarray)):
+    assert len(shape) == 1, f'{shape}'
+
+    gens = [idx]
+  elif isinstance(idx, tuple):
+    assert len(shape) >= len(idx), f'{shape} vs. {idx}'
+
+    gens = [None] * len(shape)
+    eidx = None
+    for i, tidx in enumerate(idx):
+      if isinstance(tidx, int):
+        gens[i] = (tidx,)
+      elif isinstance(tidx, slice):
+        gens[i] = range(*tidx.indices(shape[i]))
+      elif hasattr(tidx, '__iter__'):
+        gens[i] = tidx
+      elif isinstance(tidx, Ellipsis.__class__):
+        eidx = i
+        break
+
+    if eidx is not None:
+      for i in range(1, len(idx) - eidx):
+        tidx = idx[-i]
+        if isinstance(tidx, int):
+          gens[-i] = (tidx,)
+        elif isinstance(tidx, slice):
+          gens[-i] = range(*tidx.indices(shape[-i]))
+        elif hasattr(tidx, '__iter__'):
+          gens[-i] = tidx
+        elif isinstance(tidx, Ellipsis.__class__):
+          raise ValueError(f'Wrong index {idx} for shape {shape}')
+
+      for i in range(len(shape)):
+        if gens[i] is None:
+          gens[i] = range(shape[i])
+  else:
+    raise ValueError(f'Wrong index {idx} for shape {shape}')
+
+  return itertools.product(*gens)
+