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

recurvedata/utils/imports.py

@@ -0,0 +1,132 @@
+"""
+This module is a copy of the import_string function from Sentry.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Type
+
+
+class ModuleProxyCache(dict):
+    def __missing__(self, key):
+        if "." not in key:
+            return __import__(key)
+
+        module_name, class_name = key.rsplit(".", 1)
+
+        module = __import__(module_name, {}, {}, [class_name])
+        handler = getattr(module, class_name)
+
+        # We cache a NoneType for missing imports to avoid repeated lookups
+        self[key] = handler
+
+        return handler
+
+
+_cache = ModuleProxyCache()
+
+
+def import_string(path: str) -> Type:
+    """
+    Path must be module.path.ClassName
+
+    >>> cls = import_string('sentry.models.Group')
+    """
+    result = _cache[path]
+    return result
+
+
+# This is a copy of the qualname function from Apache Airflow.
+def qualname(o: object | Callable) -> str:
+    """Convert an attribute/class/function to a string importable by ``import_string``."""
+    if callable(o) and hasattr(o, "__module__") and hasattr(o, "__name__"):
+        return f"{o.__module__}.{o.__name__}"
+
+    cls = o
+
+    if not isinstance(cls, type):  # instance or class
+        cls = type(cls)
+
+    name = cls.__qualname__
+    module = cls.__module__
+
+    if module and module != "__builtin__":
+        return f"{module}.{name}"
+
+    return name
+
+
+class MockModule:
+    """
+    A placeholder object for handling failed module imports.
+    When attempting to access any attribute or call any method on this object,
+    it will raise an ImportError exception to notify the user that the module
+    was not successfully imported.
+
+    Usage example:
+    ```
+    try:
+        import some_module
+    except ImportError:
+        some_module = MockModule("some_module")
+
+    # When trying to use some_module, it will raise an import error
+    some_module.some_function()  # Will raise ImportError
+    ```
+    """
+
+    def __init__(self, module_name):
+        self.__module_name = module_name
+
+    def __getattr__(self, name):
+        """Triggered when accessing any attribute"""
+        self.__raise_import_error()
+
+    def __call__(self, *args, **kwargs):
+        """Triggered when attempting to call as a function"""
+        self.__raise_import_error()
+
+    def __getitem__(self, key):
+        """Triggered when attempting to access as dictionary or list"""
+        self.__raise_import_error()
+
+    def __raise_import_error(self):
+        """Raise standard import error"""
+        raise ImportError(
+            f"Module '{self.__module_name}' was not successfully imported. Please install the module before using it."
+        )
+
+
+class MockDecorator:
+    """
+    A decorator that raises ImportError when the decorated function is called.
+
+    This is useful for creating placeholder decorators for optional dependencies.
+    The ImportError is only raised when the decorated function is actually called,
+    not when the module is imported or the function is defined.
+
+    Usage example:
+    ```
+    try:
+        from optional_package import some_decorator
+    except ImportError:
+        some_decorator = MockDecorator("optional_package")
+
+    @some_decorator
+    def my_function():
+        pass
+
+    # The ImportError will only be raised when my_function is called, not when it's defined
+    ```
+    """
+
+    def __init__(self, module_name):
+        self.__module_name = module_name
+
+    def __call__(self, func, *args, **kwargs):
+        def wrapper(*args, **kwargs):
+            raise ImportError(
+                f"Module '{self.__module_name}' was not successfully imported. Please install the module before using it."
+            )
+
+        return wrapper

recurvedata/utils/json.py

@@ -0,0 +1,80 @@
+import datetime
+import decimal
+import json
+import uuid
+from typing import Any
+
+try:
+    import orjson
+except ImportError:
+    orjson = None
+
+
+def _json_default(obj: Any) -> str:
+    if isinstance(obj, datetime.date):
+        return obj.isoformat()
+    if isinstance(obj, datetime.timedelta):
+        return str(obj)
+    if isinstance(obj, uuid.UUID):
+        return str(obj)
+    if isinstance(obj, decimal.Decimal):
+        return str(obj)
+    raise TypeError(f"Object of type '{obj.__class__.__name__}' is not JSON serializable")
+
+
+def _orjson_default(obj: Any) -> str:
+    if isinstance(obj, datetime.timedelta):
+        return str(obj)
+    if isinstance(obj, decimal.Decimal):
+        return str(obj)
+    raise TypeError(f"Object of type '{obj.__class__.__name__}' is not JSON serializable")
+
+
+def pretty_print(v):
+    print(dumps(v, indent=2, ensure_ascii=False))
+
+
+def dumps(data: Any, **kwargs) -> str:
+    """Serialize ``data`` to JSON. Uses orjson if available."""
+
+    if orjson is None:
+        if not kwargs.get("indent", False):
+            kwargs.setdefault("separators", (",", ":"))
+        return json.dumps(data, default=_json_default, **kwargs)
+
+    # orjson does not support all the same kwargs as json.dumps
+    option = orjson.OPT_NON_STR_KEYS
+    if kwargs.pop("indent", False):
+        option |= orjson.OPT_INDENT_2
+    if kwargs.pop("sort_keys", False):
+        option |= orjson.OPT_SORT_KEYS
+
+    return orjson.dumps(data, default=_orjson_default, option=option).decode()
+
+
+def loads(data: str) -> Any:
+    """Deserialize ``data`` from JSON. Uses orjson if available."""
+
+    if orjson is None:
+        return json.loads(data)
+
+    return orjson.loads(data)
+
+
+class JSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, datetime.date):
+            return obj.isoformat()
+
+        if isinstance(obj, datetime.timedelta):
+            return str(obj)
+
+        super(JSONEncoder, self).default(obj)
+
+
+def json_dumps(data, **kwargs):
+    return json.dumps(data, cls=JSONEncoder, **kwargs)
+
+
+def json_loads(content: str, **kwargs):
+    return json.loads(content, **kwargs)

recurvedata/utils/log.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import datetime
+import inspect
+import logging
+from typing import Optional, Type, TypeVar, Union
+
+import pendulum
+
+_T = TypeVar("_T")
+
+
+# This is a copy of the airflow.utils.log.logging_mixin.LoggingMixin class
+class LoggingMixin:
+    """Convenience super-class to have a logger configured with the class name."""
+
+    _log: Optional[logging.Logger] = None
+
+    # Parent logger used by this class. It should match one of the loggers defined in the
+    # `logging_config_class`. By default, this attribute is used to create the final name of the logger, and
+    # will prefix the `_logger_name` with a separating dot.
+    _log_config_logger_name: Optional[str] = None  # noqa: UP007
+
+    _logger_name: Optional[str] = None  # noqa: UP007
+
+    @staticmethod
+    def _create_logger_name(
+        logged_class: Type[_T],
+        log_config_logger_name: str = None,
+        class_logger_name: str = None,
+    ) -> str:
+        """Generate a logger name for the given `logged_class`.
+
+        By default, this function returns the `class_logger_name` as logger name. If it is not provided,
+        the {class.__module__}.{class.__name__} is returned instead. When a `parent_logger_name` is provided,
+        it will prefix the logger name with a separating dot.
+        """
+        logger_name: str = (
+            class_logger_name if class_logger_name is not None else f"{logged_class.__module__}.{logged_class.__name__}"
+        )
+
+        if log_config_logger_name:
+            return f"{log_config_logger_name}.{logger_name}" if logger_name else log_config_logger_name
+        return logger_name
+
+    @classmethod
+    def _get_log(cls, obj: Union["LoggingMixin", Type["LoggingMixin"]], clazz: Type[_T]) -> logging.Logger:
+        if obj._log is None:
+            logger_name: str = cls._create_logger_name(
+                logged_class=clazz,
+                log_config_logger_name=obj._log_config_logger_name,
+                class_logger_name=obj._logger_name,
+            )
+            obj._log = logging.getLogger(logger_name)
+        return obj._log
+
+    @classmethod
+    def logger(cls) -> logging.Logger:
+        """Return a logger."""
+        return LoggingMixin._get_log(cls, cls)
+
+    @property
+    def log(self) -> logging.Logger:
+        """Return a logger."""
+        return LoggingMixin._get_log(self, self.__class__)
+
+
+class AwareFormatter(logging.Formatter):
+    _local_tz = pendulum.local_timezone()
+
+    def formatTime(self, record, datefmt=None):
+        # Use dateutil to get a timezone-aware datetime
+        dt = datetime.datetime.fromtimestamp(record.created, tz=self._local_tz)
+        if datefmt:
+            return dt.strftime(datefmt)
+        return dt.isoformat()
+
+
+def init_logging(
+    level=logging.INFO,
+    fmt="%(asctime)s - %(levelname)s - %(name)s - %(filename)s:%(lineno)d - [%(process)d:%(threadName)s] - %(message)s",
+):
+    logging.basicConfig(level=level, format=fmt)
+
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+
+
+def setup_loguru():
+    class InterceptHandler(logging.Handler):
+        """Intercept standard logging messages and redirect them to loguru.""" ""
+
+        def emit(self, record):
+            from loguru import logger
+
+            level: str | int
+            try:
+                level = logger.level(record.levelname).name
+            except ValueError:
+                level = record.levelno
+
+            # Find caller from where originated the logged message.
+            frame, depth = inspect.currentframe(), 0
+            while frame and (depth == 0 or frame.f_code.co_filename == logging.__file__):
+                frame = frame.f_back
+                depth += 1
+
+            logger.opt(depth=depth, exception=record.exc_info).log(level, record.getMessage())
+
+    # intercept everything at the root logger
+    logging.root.handlers = [InterceptHandler()]
+    # logging.root.setLevel(logging.INFO)
+
+    # remove every other logger's handlers
+    # and propagate to root logger
+    for name in logging.root.manager.loggerDict.keys():
+        logging.getLogger(name).handlers = []
+        logging.getLogger(name).propagate = True

recurvedata/utils/log_capture.py

@@ -0,0 +1,153 @@
+import logging
+import sys
+import threading
+import traceback
+from queue import Empty, Queue
+from typing import Optional, Protocol, Type
+
+from recurvedata.utils.log import AwareFormatter
+
+
+class MessageHandler(Protocol):
+    def __call__(self, message: str) -> None:
+        ...
+
+
+class OutputInterceptor:
+    def __init__(self, handler: MessageHandler, flush_interval_seconds: int = 5, batch_size: int = 10) -> None:
+        """
+        Initialize the OutputInterceptor object.
+
+        Args:
+            handler: The handler to call with processed messages.
+            flush_interval_seconds: Time interval (in seconds) between flushes.
+            batch_size: Number of messages to accumulate before triggering a flush.
+        """
+        self.handler = handler
+        self.flush_interval_seconds = flush_interval_seconds
+        self.batch_size = batch_size
+
+        self.queue: Queue[str] = Queue()
+        self._stop_event = threading.Event()
+        self._flusher_thread = threading.Thread(target=self._periodic_flush, daemon=True)
+        self._flusher_thread.start()
+
+        # Create a dedicated logger for internal use
+        self._logger = logging.getLogger(self.logger_name)
+
+        self._original_stdout = sys.stdout
+        self._original_stderr = sys.stderr
+
+    @property
+    def logger_name(self) -> str:
+        return f"{__name__}.{self.__class__.__name__}"
+
+    def __enter__(self) -> "OutputInterceptor":
+        """Support for context management, starts the interceptor."""
+        self._original_stdout = sys.stdout
+        self._original_stderr = sys.stderr
+        sys.stdout = self
+        sys.stderr = self
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type],
+        exc_value: Optional[BaseException],
+        exc_traceback: Optional[BaseException],
+    ) -> None:
+        """Ensure all remaining data is flushed when the context exits."""
+        if exc_type is not None:
+            tb_lines = traceback.format_exception(exc_type, exc_value, exc_traceback)
+            self._write("".join(tb_lines))
+
+        self.stop()
+
+        sys.stdout = self._original_stdout
+        sys.stderr = self._original_stderr
+
+    def write(self, s: str) -> None:
+        """
+        Write a string to the queue. This method is called when sys.stdout or sys.stderr is written to.
+        """
+        # Write to the original output stream
+        if sys.stdout is self:
+            self._original_stdout.write(s)
+        elif sys.stderr is self:
+            self._original_stderr.write(s)
+
+        self._write(s)
+
+    def flush(self) -> None:
+        """No-op flush method to maintain compatibility with sys.stdout and sys.stderr."""
+        pass
+
+    def write_log(self, s: str) -> None:
+        self._write(s + "\n")
+
+    def _write(self, s: str) -> None:
+        if isinstance(s, bytes):
+            # when an error raised, the last line of s could be bytes
+            s = s.decode("utf-8", errors="replace")
+        # if s.strip():  # Ignore empty lines
+        #     self.queue.put(s)
+        self.queue.put(s)
+
+        if self.queue.qsize() >= self.batch_size:
+            self.flush_messages()
+        elif self._stop_event.is_set():
+            # when Exception, because we Propagate the exception, the exception will write to stderr after __exit__
+            self.flush_messages()
+
+    def _periodic_flush(self) -> None:
+        """Periodically flush the queue and process each message using the processor."""
+        while not self._stop_event.is_set():
+            try:
+                self.flush_messages()
+            except Exception as e:
+                self._logger.error(f"Error during message flush: {e}", exc_info=True)
+            self._stop_event.wait(self.flush_interval_seconds)
+
+    def flush_messages(self) -> None:
+        """Flush all queued messages using the provided processor."""
+        messages = []
+        while True:
+            try:
+                messages.append(self.queue.get_nowait())
+            except Empty:
+                break
+        if messages:
+            try:
+                self.handler("".join(messages))
+            except Exception as e:
+                self._logger.error(f"Failed to process messages: {e}", exc_info=True)
+
+    def stop(self) -> None:
+        """Stop the periodic flush thread and ensure any remaining data is processed."""
+        self._stop_event.set()
+        self._flusher_thread.join()
+        self.flush_messages()  # Ensure all remaining data is processed
+
+
+class LoggingHandler(logging.Handler):
+    def __init__(self, interceptor: OutputInterceptor) -> None:
+        super().__init__()
+        self.interceptor = interceptor
+
+    def emit(self, record: logging.LogRecord) -> None:
+        # Filter out logs from the OutputInterceptor's own logger
+        if record.name == self.interceptor.logger_name:
+            return
+
+        log_entry = self.format(record)
+        self.interceptor.write_log(log_entry)
+
+
+def setup_log_handler(
+    interceptor: OutputInterceptor,
+    fmt="[%(asctime)s] - %(levelname)s - %(name)s - %(filename)s:%(lineno)d - [%(process)d:%(threadName)s] - %(message)s",
+    formatter_cls: Type[logging.Formatter] = AwareFormatter,
+):
+    handler = LoggingHandler(interceptor)
+    handler.setFormatter(formatter_cls(fmt))
+    logging.getLogger().addHandler(handler)

recurvedata/utils/mp.py

@@ -0,0 +1,178 @@
+import logging
+import subprocess
+import sys
+import time
+from multiprocessing import Process
+from multiprocessing.queues import Queue
+from queue import Empty, Full
+from typing import Any, Optional, Union
+
+logger = logging.getLogger(__name__)
+
+
+def get_qsize(queue: Queue) -> Optional[int]:
+    if sys.platform.lower() == "darwin":
+        # queue.qsize() Raises NotImplementedError on Mac OSX because of broken sem_getvalue()
+        return None
+    return queue.qsize()
+
+
+def safe_join_subprocesses(workers: list[Process], result_queue: Queue) -> list[Any]:
+    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: bool = True) -> bool:
+    for p in workers:
+        if p.is_alive():
+            continue
+        if p.exitcode != 0:
+            if log:
+                logger.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():
+            logger.info(f"start terminate process {p.pid}")
+            p.terminate()
+            logger.info(f"finish terminate process {p.pid}")
+
+
+def master_safe_put_queue(
+    queue: Queue,
+    obj: Any,
+    workers: list[Process],
+    block: bool = True,
+    timeout: Optional[int] = None,
+) -> Optional[bool]:
+    """
+    A scenario where queue.put is called involves the master putting data into the queue, while workers consume the data.
+     Under the default settings of timeout=None and block=True,
+     if queue.maxsize is small and workers encounter errors that prevent them from consuming data promptly,
+     this can cause the master to become stuck when calling queue.put.
+    This master_safe_put_queue function can address this issue.
+    When operating under timeout=None and block=True,
+     it uses a small timeout (10 seconds) and continuously attempts queue.put(timeout=10) in a loop.
+     If queue.put becomes stuck for 10 seconds, it raises a queue.Full error.
+
+    Args:
+        queue: multiprocessing.Queue
+        obj: the object to be placed into the queue
+        workers: subprocesses
+        block: whether to block when the queue has no free slots
+        timeout: the timeout for queue.put
+
+    Returns:
+        If there are any worker processes that have exited abnormally, it returns True,
+        indicating that an abnormal worker exit caused the master's queue.put to become stuck.
+        If all workers are functioning normally, it indicates that the workers are simply consuming data slowly,
+        and it will continue to call queue.put(timeout=10) in a loop.
+        In other cases, it behaves the same as queue.put.
+    """
+    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]:
+    """
+    wait and read the sub workers' result from result_queue,
+    exit when
+        1) one sub worker fail, or
+        2) all sub workers success
+
+    Args:
+        workers: sub progresses
+        result_queue: queue which sub progresses put result into
+
+    Returns:
+        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],
+    return_output=False,
+    _logger: logging.Logger = logger,
+    **kwargs,
+) -> Optional[str]:
+    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True, **kwargs)
+    logger.info(f"started sub process: {cmd}, pid: {p.pid}")
+    lines: list[str] = []
+    for raw_line in iter(p.stdout.readline, ""):
+        line = raw_line.rstrip()
+        _logger.info(line)
+        if return_output:
+            lines.append(line)
+    p.wait()
+    logger.info("sub process exited with return code %s", p.returncode)
+    if p.returncode:
+        raise subprocess.CalledProcessError(p.returncode, p.args)
+    return "\n".join(lines)
+
+
+def robust_run_subprocess(
+    cmd: Union[str, list],
+    _logger: logging.Logger = logger,
+    **kwargs,
+) -> tuple[str, int]:
+    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True, **kwargs)
+    logger.info(f"started sub process: {cmd}, pid: {p.pid}")
+    lines: list[str] = []
+    for raw_line in iter(p.stdout.readline, ""):
+        line = raw_line.rstrip()
+        _logger.info(line)
+        lines.append(line)
+    p.wait()
+    logger.info("sub process exited with return code %s", p.returncode)
+    output = "\n".join(lines)
+    return output, p.returncode