recurvedata-lib 0.1.487__py2.py3-none-any.whl

recurvedata/pigeon/utils/fs.py

@@ -0,0 +1,266 @@
+import bz2
+import contextlib
+import datetime
+import glob
+import gzip
+import logging
+import os
+import shutil
+import subprocess
+import tempfile
+import zipfile
+from itertools import islice
+
+from recurvedata.pigeon.utils import ensure_list
+
+
+def new_tempfile(suffix="", prefix=None, dir=None):
+    ts = datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
+    suffix = "{}_{}".format(ts, suffix)
+    kwargs = {"suffix": suffix, "dir": dir}
+    if prefix:
+        kwargs["prefix"] = prefix
+    _, filename = tempfile.mkstemp(**kwargs)
+    return filename
+
+
+class new_stagefile_factory:
+    def __init__(self, directory):
+        if not os.path.isabs(directory):
+            directory = os.path.join("/tmp", directory)
+        self.directory = directory
+
+    def __call__(self, name):
+        os.makedirs(self.directory, exist_ok=True)
+        return os.path.join(self.directory, name)
+
+
+def merge_files(files, filename=None, num_skip_lines=0, delete=True):
+    """Concat multiple files into one file.
+
+    :param files: source file names
+    :param filename: target filename, will create a tempfile if not provided
+    :param num_skip_lines: skip n lines before merge into target file
+    :param delete: delete source files after been merged
+    :return: the target filename
+    """
+    if filename is None:
+        _, filename = tempfile.mkstemp()
+
+    if not num_skip_lines:
+        if len(files) == 1 and delete:
+            # just rename
+            os.rename(files[0], filename)
+        else:
+            # merge by `cat` for better performance
+            cmd = f'cat {" ".join(files)} > {filename}'
+            _run_command(cmd)
+    else:
+        with open(filename, "wb") as fout:
+            for f in files:
+                with open(f, "rb") as fin:
+                    for _ in range(num_skip_lines):
+                        fin.readline()
+                    shutil.copyfileobj(fin, fout)
+
+    if delete:
+        remove_files_safely(files)
+
+    return filename
+
+
+def skip_lines(infile, lines, inplace=False):
+    tmpfile = new_tempfile()
+    with open(infile, "rb") as fin, open(tmpfile, "wb") as fout:
+        # skip the first n lines
+        for _ in range(lines):
+            fin.readline()
+
+        # copy the rest to another file
+        shutil.copyfileobj(fin, fout)
+
+    if inplace:
+        os.rename(tmpfile, infile)
+        return infile
+    return tmpfile
+
+
+def read_lines(filename, start_line, lines_num=1):
+    with open(filename) as f:
+        for line in islice(f, start_line, start_line + lines_num):
+            yield line
+
+
+def is_file_empty(filename):
+    """Detect file is empty or not, the non-exists file is considered as empty"""
+    try:
+        return os.stat(filename).st_size == 0
+    except FileNotFoundError:
+        return True
+
+
+def remove_files(files):
+    for f in ensure_list(files):
+        os.unlink(f)
+
+
+def remove_files_safely(files):
+    with contextlib.suppress(OSError, TypeError, ValueError):
+        remove_files(files)
+
+
+def remove_files_by_pattern(pattern):
+    files = glob.glob(pattern)
+    logging.info("files to be deleted: %s", str(files))
+    remove_files_safely(files)
+
+
+def remove_folder_safely(folder):
+    if not os.path.exists(folder):
+        return
+    shutil.rmtree(folder, ignore_errors=True)
+
+
+def gzip_compress(filename, target_filename=None, using_cmd=False):
+    """Compress a file using gzip
+    :param filename: the path of input file
+    :param target_filename: the path of output file, a temporary filename will be made otherwise
+    :param using_cmd: use the gzip command line instead of Python GzipFile to speedup
+    :return: the target_filename
+    """
+    if target_filename is None:
+        target_filename = new_tempfile(suffix=".gz")
+
+    if using_cmd:
+        _run_command(f"gzip {filename} -c > {target_filename}")
+        return target_filename
+
+    with open(filename, "rb") as f_in, gzip.GzipFile(target_filename, "wb") as f_out:
+        shutil.copyfileobj(f_in, f_out)
+    return target_filename
+
+
+def gzip_decompress(filename, target_filename=None, using_cmd=False):
+    """Decompress a gzip file
+    :param filename: the path of the gzip file
+    :param target_filename: the path of output file, a temporary filename will be made otherwise
+    :param using_cmd: use the gzip command line instead of Python GzipFile to speedup
+    :return: the target_filename
+    """
+    if target_filename is None:
+        target_filename = new_tempfile()
+
+    if using_cmd:
+        _run_command(f"gzip -d {filename} -c > {target_filename}")
+        return target_filename
+
+    with gzip.GzipFile(filename, "rb") as f_in, open(target_filename, "wb") as f_out:
+        shutil.copyfileobj(f_in, f_out)
+    return target_filename
+
+
+def bzip2_compress(filename, target_filename=None, using_cmd=False):
+    """Compress a file using bzip2
+    :param filename: the path of input file
+    :param target_filename: the path of output file, a temporary filename will be made otherwise
+    :param using_cmd: use the bzip2 command line instead of Python BZ2File to speedup
+    :return: the target_filename
+    """
+    if target_filename is None:
+        target_filename = new_tempfile(suffix=".bz2")
+
+    if using_cmd:
+        _run_command(f"bzip2 {filename} -c > {target_filename}")
+        return target_filename
+
+    with open(filename, "rb") as f_in, bz2.BZ2File(target_filename, "wb") as f_out:
+        shutil.copyfileobj(f_in, f_out)
+    return target_filename
+
+
+def bzip2_decompress(filename, target_filename=None, using_cmd=False):
+    """Decompress a bzip2 file
+    :param filename: the path of the bzip2 file
+    :param target_filename: the path of output file, a temporary filename will be made otherwise
+    :param using_cmd: use the gzip command line instead of Python BZ2File to speedup
+    :return: the target_filename
+    """
+    if target_filename is None:
+        target_filename = new_tempfile()
+
+    if using_cmd:
+        _run_command(f"bzip2 -d {filename} -c > {target_filename}")
+        return target_filename
+
+    with bz2.BZ2File(filename, "rb") as f_in, open(target_filename, "wb") as f_out:
+        shutil.copyfileobj(f_in, f_out)
+    return target_filename
+
+
+def zip_compress(filename, target_filename=None, using_cmd=False, arcname=None):
+    """Compress a file using zip
+    :param filename: the path of input file
+    :param target_filename: the path of output file, a temporary filename will be made otherwise
+    :param using_cmd: use the zip command line instead of Python ZipFile to speedup
+    :param arcname: filename in the archive file, only supported with using_cmd=False
+    :return: the target_filename
+    """
+    if target_filename is None:
+        target_filename = new_tempfile(suffix=".zip")
+
+    directory, basename = os.path.split(filename.rstrip("/"))
+
+    if using_cmd:
+        # 先删除生成的临时文件，只使用生成的文件名，要不然会报错
+        # zip warning: missing end signature--probably not a zip file (did you
+        # zip warning: remember to use binary mode when you transferred it?)
+        # zip warning: (if you are trying to read a damaged archive try -F)
+        remove_files_safely(target_filename)
+        if arcname is not None:
+            logging.warning("arcname is not supported while using cmd")
+        _run_command(f"cd {directory} && zip -r {target_filename} {basename}")
+        return target_filename
+
+    with zipfile.ZipFile(target_filename, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+        zf.write(filename, arcname=arcname or basename)
+    return target_filename
+
+
+def zip_decompress(filename, target_directory=None, using_cmd=False):
+    """Decompress a .zip file
+    :param filename: the path of input file
+    :param target_directory: the path of output directory, a temporary directory will be made otherwise
+    :param using_cmd: use the unzip command line instead of Python ZipFile to speedup
+    :return: the output directory
+    """
+    if not target_directory:
+        target_directory = tempfile.mkdtemp()
+
+    if using_cmd:
+        _run_command(f"unzip {filename} -d {target_directory}")
+        return target_directory
+
+    with zipfile.ZipFile(filename, "r") as zf:
+        zf.extractall(target_directory)
+    return target_directory
+
+
+@contextlib.contextmanager
+def ensure_remove(filename):
+    try:
+        yield filename
+    finally:
+        remove_files_safely(filename)
+
+
+def schema_filename(base):
+    return f"{base}.schema"
+
+
+def exists(path):
+    return os.path.exists(path)
+
+
+def _run_command(cmd):
+    logging.info(cmd)
+    subprocess.check_output(cmd, shell=True)

recurvedata/pigeon/utils/json.py

@@ -0,0 +1,44 @@
+import datetime
+import decimal
+import json
+
+
+class JSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, decimal.Decimal):
+            return float(obj)
+
+        if isinstance(obj, (datetime.date, datetime.datetime)):
+            return obj.isoformat()
+
+        if isinstance(obj, datetime.timedelta):
+            return str(obj)
+
+        return super().default(obj)
+
+
+def json_dumps(obj, **kwargs):
+    kwargs.setdefault('cls', JSONEncoder)
+    return json.dumps(obj, **kwargs)
+
+
+def dump_json(obj, fp=None, **kwargs):
+    kwargs.setdefault('indent', 4)
+    kwargs.setdefault('ensure_ascii', False)
+    kwargs.setdefault('sort_keys', True)
+    kwargs.setdefault('cls', JSONEncoder)
+    if fp is None:
+        return json.dumps(obj, **kwargs)
+    else:
+        if isinstance(fp, str):
+            with open(fp, 'w') as fp:
+                return json.dump(obj, fp, **kwargs)
+        return json.dump(obj, fp, **kwargs)
+
+
+def load_json(fp, **kwargs):
+    if isinstance(fp, str):
+        with open(fp, 'r') as fp:
+            return json.load(fp, **kwargs)
+    else:
+        return json.load(fp, **kwargs)

recurvedata/pigeon/utils/keyed_tuple.py

@@ -0,0 +1,85 @@
+"""
+This file is taken from SQLAlchemy
+"""
+
+
+class AbstractKeyedTuple(tuple):
+    __slots__ = ()
+
+    def keys(self):
+        """Return a list of string key names for this :class:`.KeyedTuple`.
+
+        .. seealso::
+
+            :attr:`.KeyedTuple._fields`
+
+        """
+
+        return list(self._fields)
+
+
+class KeyedTuple(AbstractKeyedTuple):
+    """``tuple`` subclass that adds labeled names.
+
+    E.g.::
+
+        >>> k = KeyedTuple([1, 2, 3], labels=["one", "two", "three"])
+        >>> k.one
+        1
+        >>> k.two
+        2
+
+    Result rows returned by :class:`.Query` that contain multiple
+    ORM entities and/or column expressions make use of this
+    class to return rows.
+
+    The :class:`.KeyedTuple` exhibits similar behavior to the
+    ``collections.namedtuple()`` construct provided in the Python
+    standard library, however is architected very differently.
+    Unlike ``collections.namedtuple()``, :class:`.KeyedTuple` is
+    does not rely on creation of custom subtypes in order to represent
+    a new series of keys, instead each :class:`.KeyedTuple` instance
+    receives its list of keys in place.   The subtype approach
+    of ``collections.namedtuple()`` introduces significant complexity
+    and performance overhead, which is not necessary for the
+    :class:`.Query` object's use case.
+
+    .. seealso::
+
+        :ref:`ormtutorial_querying`
+
+    """
+
+    def __new__(cls, vals, labels=None):
+        t = tuple.__new__(cls, vals)
+        if labels:
+            t.__dict__.update(zip(labels, vals))
+        else:
+            labels = []
+        t.__dict__["_labels"] = labels
+        return t
+
+    @property
+    def _fields(self):
+        """Return a tuple of string key names for this :class:`.KeyedTuple`.
+
+        This method provides compatibility with ``collections.namedtuple()``.
+
+        .. seealso::
+
+            :meth:`.KeyedTuple.keys`
+
+        """
+        return tuple([l for l in self._labels if l is not None])
+
+    def __setattr__(self, key, value):
+        raise AttributeError("Can't set attribute: %s" % key)
+
+    def _asdict(self):
+        """Return the contents of this :class:`.KeyedTuple` as a dictionary.
+
+        This method provides compatibility with ``collections.namedtuple()``,
+        with the exception that the dictionary returned is **not** ordered.
+
+        """
+        return {key: self.__dict__[key] for key in self.keys()}

recurvedata/pigeon/utils/mp.py

@@ -0,0 +1,156 @@
+import logging
+import threading
+import time
+from multiprocessing import Process
+from multiprocessing.queues import Queue
+from queue import Empty, Full
+from subprocess import PIPE, STDOUT, CalledProcessError, Popen
+from typing import Any, List, Optional, Tuple, Union
+
+
+def safe_join_subprocesses(workers, result_queue):
+    result = []
+    live_workers = list(workers)
+    while live_workers:
+        try:
+            while 1:
+                result.append(result_queue.get(False))
+        except Empty:
+            pass
+
+        time.sleep(0.5)  # Give tasks a chance to put more data in
+        if not result_queue.empty():
+            continue
+        live_workers = [p for p in live_workers if p.is_alive()]
+    return result
+
+
+def has_process_fail(workers: List[Process], log=True):
+    for p in workers:
+        if p.is_alive():
+            continue
+        if p.exitcode != 0:
+            if log:
+                logging.info(f"found process {p.pid} fail, exitcode {p.exitcode}")
+            return True
+    return False
+
+
+def terminate_processes(workers: List[Process]):
+    for p in workers:
+        if p.is_alive():
+            logging.info(f"start terminate process {p.pid}")
+            p.terminate()
+            logging.info(f"finish terminate process {p.pid}")
+
+
+def master_safe_put_queue(
+    queue: Queue, obj: Any, workers: List[Process], block=True, timeout: Optional[int] = None
+) -> Optional[bool]:
+    """
+    一种调用 queue.put 的场景，是 master put 数据，worker 消费数据.
+    在默认的 timeout=None, block=True 下，
+    如果 queue.maxsize 较小，且 workers 遇到了报错，没法及时消费，
+    就会导致 master 在调用 queue.put 的时候卡住。
+    master_safe_put_queue 可以解决这个问题，
+    当 timeout=None, block=True 的情况下，
+    会用一个较小的 timeout（10s），死循环不断尝试 queue.put(timeout=10)，
+    当 queue.put 卡住达到 10s 的时候，会报错 queue.Full，
+    这时候检查 workers 是否有异常退出的进程，
+        如果 workers 有异常退出的进程，则返回 True, 表示 worker 有异常退出导致 master queue.put 卡住;
+        如果 workers 都正常，则表示确实是 worker 消费速度较慢，重新调用 queue.put(timeout=10) 继续死循环
+
+    其他情况下与 queue.put 一致
+
+    :param queue: queue
+    :param obj: the obj to put into queue
+    :param workers: sub processes
+    :param block: should block when queue has no free slot
+    :param timeout: queue.put's timeout
+    :return: True 表示 workers 有异常退出导致 master queue.put 卡住；否则返回 None
+    """
+    if timeout is None and block:
+        while True:
+            try:
+                return queue.put(obj, timeout=10)
+            except Full:
+                if has_process_fail(workers):
+                    return True
+    else:
+        return queue.put(obj, block=block, timeout=timeout)
+
+
+def safe_join_subprocesses_early_stop(workers: List[Process], result_queue: Queue) -> Tuple[List, bool]:
+    """
+    this function wait and read the sub workers' result from result_queue,
+    exit when
+        1) one sub worker fail
+        or
+        2) all sub workers success
+    :param workers: sub progresses
+    :param result_queue: queue which sub progresses put result into
+    :return: result got from sub workers, and early_stop flag
+    """
+    result = []
+    early_stop = False
+    live_workers = list(workers)
+    last_check_early_stop_time = time.time()
+    while live_workers:
+        try:
+            while 1:
+                result.append(result_queue.get(False))
+
+                if time.time() - last_check_early_stop_time > 10:
+                    if has_process_fail(live_workers):
+                        early_stop = True
+                        return result, early_stop
+                    last_check_early_stop_time = time.time()
+
+        except Empty:
+            pass
+
+        time.sleep(0.5)  # Give tasks a chance to put more data in
+        if not result_queue.empty():
+            continue
+
+        if has_process_fail(live_workers):
+            early_stop = True
+            return result, early_stop
+        last_check_early_stop_time = time.time()
+        live_workers = [p for p in live_workers if p.is_alive()]
+    return result, early_stop
+
+
+def run_subprocess(cmd: Union[str, List], stdout=PIPE, stderr=STDOUT, return_output=False, **kwargs) -> Optional[str]:
+    p = Popen(cmd, stdout=stdout, stderr=stderr, **kwargs)
+    logging.info(f"started sub process: {cmd}, pid: {p.pid}")
+    lines: List[str] = []
+    for raw_line in iter(p.stdout.readline, b""):
+        line = raw_line.decode("utf8").rstrip()
+        logging.info(line)
+        if return_output:
+            lines.append(line)
+    p.wait()
+    logging.info("sub process exited with return code %s", p.returncode)
+    if p.returncode:
+        raise CalledProcessError(p.returncode, p.args)
+    return "\n".join(lines)
+
+
+class PropagatingThread(threading.Thread):
+    def run(self):
+        self.exc = None
+        try:
+            if hasattr(self, "_Thread__target"):
+                # Thread uses name mangling prior to Python 3.
+                self.ret = self._Thread__target(*self._Thread__args, **self._Thread__kwargs)
+            else:
+                self.ret = self._target(*self._args, **self._kwargs)
+        except BaseException as e:
+            self.exc = e
+
+    def join(self, timeout=None):
+        super().join(timeout)
+        if self.exc:
+            raise self.exc
+        return self.ret