MeUtils 2024.3.4.13.4.45__py3-none-any.whl → 2025.1.16.17.15.52__py3-none-any.whl

meutils/decorators/common.py

@@ -21,7 +21,7 @@ import wrapt
 from loguru import logger
 from tqdm.auto import tqdm
 from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
-from contextlib import contextmanager
+from contextlib import contextmanager, asynccontextmanager
 from typing import *
 # ME
 from meutils.decorators.decorator import decorator
@@ -77,7 +77,7 @@ def return2log(func, sink=sys.stderr, logkwargs=None, *args, **kwargs):
 
 
 @contextmanager
-def timer(task="Task"):
+def timer(task="Task", notice_fn: Optional[Callable] = None):
     """https://www.kaggle.com/lopuhin/mercari-golf-0-3875-cv-in-75-loc-1900-s
         # 其他装饰器可学习这种写法
         with timer() as t:
@@ -88,13 +88,50 @@ def timer(task="Task"):
             print('sleeping')
             time.sleep(6)
             return 6
+
+
+        from meutils.notice import feishu
+        feishu.send_message
     """
 
     logger.info(f"{task} started")
     s = time.perf_counter()
     yield
     e = time.perf_counter()
-    logger.info(f"{task} done in {e - s:.3f} s")
+    msg = f"{task} done in {e - s:.3f} s"
+    logger.info(msg)  # feishu
+
+    # notice
+    if notice_fn: notice_fn(msg)
+
+
+@asynccontextmanager
+async def async_timer(task="Task", notice_fn: Optional[Callable] = None):
+    """
+    Async context manager for timing tasks.
+
+    Usage:
+    ```python
+    async with async_timer() as t:
+        await asyncio.sleep(3)
+
+    @async_timer()
+    async def f():
+        print('sleeping')
+        await asyncio.sleep(6)
+        return 6
+    ```
+    """
+    logger.info(f"{task} started")
+    s = time.perf_counter()
+    yield
+    e = time.perf_counter()
+    msg = f"{task} done in {e - s:.3f} s"
+    logger.info(msg)  # feishu
+
+    # notice
+    if notice_fn:
+        await notice_fn(msg)
 
 
 @contextmanager
@@ -170,6 +207,7 @@ def background_task(func, max_workers=1, *args, **kwargs):
     # pool.shutdown(wait=False)  # 不等待
     # pool.shutdown(wait=True)  # 等待
 
+    # with ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix='🐶') as pool: # 失去异步效果
     pool = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix='🐶')
     future = pool.submit(func, *args, **kwargs)  # pool.map(fun4, ips)
     future.add_done_callback(
@@ -177,6 +215,8 @@ def background_task(func, max_workers=1, *args, **kwargs):
     )
     # future.add_done_callback()
 
+    # pool.shutdown(wait=False)  # 不等待
+
     return future.running()  # future.done()
 
 
@@ -187,7 +227,7 @@ background = background_task
 def background_task_plus(func, *args, **kwargs):
     pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix='🐶')
     future = pool.submit(func, *args, **kwargs)  # pool.map(fun4, ips)
-    future.add_done_callback(lambda x: logger.error(future.exception()) if future.exception() else None)
+    future.add_done_callback(lambda x: logger.error(f"{future.exception()}") if future.exception() else None)
     # 详细错误 traceback.format_exc()
     return future.running()  # future.done()
 
@@ -283,6 +323,43 @@ def requires_dependencies(
     return decorator
 
 
+def ratelimit(
+        name: str = "ratelimit",
+        weight: int = 1,
+        rates: Optional[List[Any]] = None,
+):
+    """ https://github.com/vutran1710/PyrateLimiter
+    :param name:
+    :param weight:
+    :param rates:
+    :return:
+    """
+    from pyrate_limiter import Duration, Rate, Limiter, BucketFullException
+
+    rates = rates or [
+        Rate(1, Duration.SECOND),
+        Rate(2, Duration.SECOND * 15),
+        Rate(2 ** 2, Duration.MINUTE),
+        Rate(2 ** 4, Duration.HOUR),
+        Rate(2 ** 8, Duration.DAY)
+    ]
+    limiter = Limiter(rates)
+
+    def name2weight(*args, **kwargs):  # 动态限频 limiter.try_acquire(self, name: str, weight: int = 1)
+        return name, weight
+
+    limiter_decorator = limiter.as_decorator()(name2weight)
+
+    @wrapt.decorator
+    def wrapper(wrapped, instance, args, kwargs):
+        try:
+            return limiter_decorator(wrapped)(*args, **kwargs)
+        except BucketFullException as e:
+            return e.meta_info
+
+    return wrapper
+
+
 def ratelimiter(limit_value: str = '3/1', callback=None):
     from ratelimiter import RateLimiter
 
@@ -321,6 +398,8 @@ def limit(limit_value: str = '1/second', key_func: Callable[..., str] = None, er
     return inner
 
 
+
+
 if __name__ == '__main__':
     import time
 
@@ -398,7 +477,7 @@ if __name__ == '__main__':
     #     pass
 
     def f():
-        return  1/0
+        return 1 / 0
         # return 1
 
 

meutils/decorators/contextmanagers.py

@@ -10,9 +10,7 @@
 
 from meutils.pipe import *
 
-from contextlib import contextmanager
-
-from urllib3 import HTTPResponse
+from contextlib import contextmanager, asynccontextmanager
 
 
 @contextmanager
@@ -37,18 +35,31 @@ def timer(task="Task"):
 
 
 @contextmanager
-def error_catcher(task, is_trace=False):
+def try_catcher(task="Task", fallback: Callable = None, is_trace: bool = False):
     try:
         yield
     except Exception as e:
         error = traceback.format_exc() if is_trace else e
         logger.error(f"{task}: {error}")
+        if fallback:
+            yield fallback()
 
 
-@contextmanager
-def try_catcher(task, is_trace=False):
+@asynccontextmanager
+async def atry_catcher(task="Task", fallback: Callable = None, is_trace: bool = False):
     try:
         yield
     except Exception as e:
         error = traceback.format_exc() if is_trace else e
         logger.error(f"{task}: {error}")
+        if fallback:
+            yield await fallback()
+
+
+if __name__ == '__main__':
+    async def f():
+        return 1/0
+
+
+    with try_catcher("test"):
+        arun(f())

meutils/decorators/fastapi_decorator.py

@@ -34,17 +34,91 @@ def catch_exceptions(func=None, *, exception_type=Exception, status_code=404):
     return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
 
 
+def catch_exceptions_(func):
+    @wraps(func)
+    async def wrapper(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except CustomException as e:
+            return JSONResponse(
+                status_code=418,
+                content={"message": f"An error occurred: {str(e)}"},
+            )
+        except Exception as e:
+            return JSONResponse(
+                status_code=500,
+                content={"message": "An unexpected error occurred."}
+            )
+
+    return wrapper
+
+
+# 然后在你的路由中使用装饰器：
+# @app.get("/items/{item_id}")
+# @catch_exceptions
+# async def read_item(item_id: int):
+#     # 你的业务逻辑
+#     pass
+
+
 if __name__ == '__main__':
     from meutils.serving.fastapi import App
+    from fastapi import Depends
 
     app = App()
 
 
+    async def catch_exceptions(exception_type=Exception, status_code=404):
+        try:
+            yield
+        except exception_type as e:
+            raise HTTPException(status_code=status_code, detail=f"Error: {e}")
+
+
+    from fastapi import FastAPI, HTTPException, Depends, Request
+    from starlette.responses import JSONResponse
+
+
+
+    async def catch_exceptions(request: Request, call_next):
+        try:
+            response = await call_next(request)
+            return response
+        except IOError as e:
+            return JSONResponse(status_code=500, content={"detail": f"IOError: {e}"})
+
+
+    @app.middleware("http")
+    async def catch_exceptions_middleware(request: Request, call_next):
+        return await catch_exceptions(request, call_next)
+
+
     @app.get("/test")
-    @catch_exceptions(exception_type=IOError, status_code=500)
-    def f(file_id: str):
-        logger.debug(file_id)
+    def test_endpoint(file_id: str):
+        # Your endpoint logic here
         raise Exception("test")
+        # return {"file_id": file_id}
+
+
+
+    # @app.get("/test_plus")
+    # def test_endpoint(file_id: str, _=Depends(catch_exceptions(exception_type=IOError, status_code=500))):
+    #     # Your endpoint logic here
+    #     return {"file_id": file_id}
+
+
+    # @app.get("/test")
+    # @catch_exceptions(exception_type=IOError, status_code=500)
+    # def f(file_id: str):
+    #     logger.debug(file_id)
+    #     # raise Exception("test")
+    #
+    #
+    # @app.get("/{file_id}")
+    # @catch_exceptions(status_code=1)
+    # def f(file_id: str):
+    #     logger.debug(file_id)
+    #     raise Exception("test")
 
 
     app.run(port=9000)

meutils/decorators/polling.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Project      : AI.  @by PyCharm
+# @File         : polling
+# @Time         : 2024/12/26 14:08
+# @Author       : betterme
+# @WeChat       : meutils
+# @Software     : PyCharm
+# @Description  :
+import time
+
+from meutils.pipe import *
+
+
+def poll(n: int = 30, sleep_fn: Optional[Callable] = None):
+    if sleep_fn is None:
+        sleep_fn = lambda i: max(n / (i + 1), 1)  # 预估大概 2~3 *n
+
+    def decorator(fn):
+        is_coroutine = inspect.iscoroutinefunction(fn)
+
+        @wraps(fn)
+        async def wrapper(*args, **kwargs):
+            for i in range(n):
+                await asyncio.sleep(sleep_fn(i))
+                # 根据函数类型选择调用方式
+                result = await fn(*args, **kwargs) if is_coroutine else fn(*args, **kwargs)
+                if result:  # 跳出
+                    return result
+            return None
+
+        return wrapper
+
+    return decorator
+
+
+if __name__ == '__main__':
+    @poll(n=5)
+    async def check_something():
+        # 你的检查逻辑
+
+        logger.debug(time.ctime())
+        return {"status": "ok"}
+
+
+    arun(check_something())

meutils/decorators/retry.py

@@ -9,29 +9,89 @@
 # @Description  :
 
 from meutils.pipe import *
-from tenacity import retry, wait_fixed, stop_after_attempt, wait_exponential, before_sleep_log
-from tenacity import retry_if_result, retry_if_exception_type  # 重试策略
+from tenacity import retry, stop_after_attempt, before_sleep_log
+from tenacity import wait_fixed, wait_exponential, wait_incrementing, wait_exponential_jitter  # Wait strategy
+from tenacity import retry_if_result, RetryCallState  # 重试策略
+from tenacity import retry_if_exception, retry_if_exception_type, retry_if_not_exception_type
 
+# 不兼容异步
+from meutils.notice.feishu import send_message
 
-@decorator
-def retrying(func, max_retries=2, min_seconds=1, max_seconds=32,
-             retry_error_callback: Callable = logger.exception,
-             predicate: Callable[[Any], bool] = lambda x: False,
-             *args, **kwargs):
+
+class IgnoredRetryException(Exception):
+    pass
+
+
+def default_retry_error_callback(retry_state: RetryCallState, title: Optional[str] = None):
+    """return the result of the last call attempt"""
+    # logger.debug(f"最后一次重试仍然失败调用的函数: {retry_state.outcome}")
+    logger.debug(f"最后一次重试仍然失败调用的函数: {retry_state}")
+
+    send_message(f"""{retry_state}""", title)
+
+    # logger.debug(retry_state.outcome.result())
+
+    exc = retry_state.outcome.exception()  # 最后抛出错误
+    if exc:
+        raise exc  # 最后返回结果
+
+
+def not_in_status_codes(exception):
+    """451不触发重试"""
+    # logger.debug(exception)
+
+    if isinstance(exception, httpx.HTTPStatusError):
+        # logger.debug(f"status_code: {exception.response.status_code}")
+
+        status_code = exception.response.status_code
+        return status_code not in {451}
+    return True
+
+
+def retrying(
+        max_retries=2,
+        exp_base=2,
+        min: int = 0,
+        max: int = 100000,
+        retry_error_callback: Optional[Callable[[RetryCallState], Any]] = None,
+        predicate: Callable[[Any], bool] = lambda r: False,
+        title: Optional[str] = None,
+        ignored_exception_types: Optional[typing.Union[
+            typing.Type[BaseException],
+            typing.Tuple[typing.Type[BaseException], ...],
+        ]] = None,  # 该错误类型不重试
+):
     import logging
 
     logger = logging.getLogger()
 
+    retry_error_callback = retry_error_callback or partial(default_retry_error_callback, title=title)
+
+    # 抛错或者返回结果判断为True重试 就重试
+    retry_reasons = (
+            (
+                    retry_if_exception_type()
+                    & retry_if_exception(not_in_status_codes)
+            )
+            | retry_if_result(predicate)
+    )
+    if ignored_exception_types:
+        retry_reasons = retry_reasons & retry_if_not_exception_type(ignored_exception_types)
+
     _retry_decorator = retry(
         reraise=True,
         stop=stop_after_attempt(max_retries),
-        wait=wait_exponential(multiplier=1, min=min_seconds, max=max_seconds),
-        retry=retry_if_exception_type() | retry_if_result(predicate),  # 抛错或者返回结果判断为True重试 就重试
-        before_sleep=before_sleep_log(logger, 30),
+        wait=wait_exponential(multiplier=1.0, exp_base=exp_base, min=min, max=max),  # max=max_seconds
+        retry=retry_reasons,
+        # before_sleep=before_sleep_log(logger, 30),
         retry_error_callback=retry_error_callback,
     )
 
-    return _retry_decorator(func)(*args, **kwargs)
+    @wrapt.decorator
+    def wrapper(wrapped, instance, args, kwargs):
+        return _retry_decorator(wrapped)(*args, **kwargs)
+
+    return wrapper
 
 
 def create_retry_decorator() -> Callable[[Any], Any]:  # todo: Retrying
@@ -84,25 +144,89 @@ def wait_retry(n=3):
 
 
 if __name__ == '__main__':
-    s = time.time()  # 1616145296
-    print(s)
-    e1 = s + 10
-    e2 = e1 + 10
+    # s = time.time()  # 1616145296
+    # print(s)
+    # e1 = s + 10
+    # e2 = e1 + 10
+    #
+    #
+    # @wait_retry(5)
+    # def f(e):
+    #     return time.time() > e  # 变的
+    #
+    #
+    # def run(e):
+    #     f(e)
+    #     print(f"task {e}")
+    #
+    #
+    # # for e in [e2, e1]:
+    # #     print(run(e))
+    # #
+    # # print("耗时", time.time() - s)
+    #
+    # [e1, e2, 1000000000000] | xProcessPoolExecutor(run, 2)
+
+    # class retry_state:
+    #     attempt_number = 0
+    #
+    #
+    # for i in range(1, 10):
+    #     retry_state.attempt_number = i
+    #     # print(wait_incrementing(100, -10)(retry_state))
+    #     print(wait_exponential(1, exp_base=2)(retry_state))
+
+    # @retrying(3)
+    # async def f():
+    #     logger.debug("retry")
+    #     pass
+    #     # 1 / 0
+    #     raise Exception('xx')
+
+    # try:
+    #     f()
+    #
+    # except Exception as e:
+    #     print(e)
+    #     print(type(e))
 
+    # @retrying(3)
+    # def f():
+    #     # async def f():
+    #     logger.debug("retry")
+    #     for i in range(10):
+    #         yield ""
+    #         if i == 8:
+    #             raise Exception('xx')
 
-    @wait_retry(5)
-    def f(e):
-        return time.time() > e  # 变的
+    # arun(f())
 
+    # f()
 
-    def run(e):
-        f(e)
-        print(f"task {e}")
+    # 示例使用
+    class CustomError1(Exception):
+        pass
 
 
-    # for e in [e2, e1]:
-    #     print(run(e))
-    #
-    # print("耗时", time.time() - s)
+    class CustomError2(Exception):
+        pass
+
+
+    @retrying(
+        max_retries=3,
+        ignored_exception_types=(CustomError1, IgnoredRetryException)
+    )
+    def my_function():
+        logger.debug(1)
+        # 模拟抛出异常
+        raise IgnoredRetryException("This is a custom error")
+
+
+    def main():
+        try:
+            my_function()
+        except Exception as e:
+            print(f"Operation failed after all retries: {e}")
+
 
-    [e1, e2, 1000000000000] | xProcessPoolExecutor(run, 2)
+    main()

meutils/fastapi_utils/__init__.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Project      : AI.  @by PyCharm
+# @File         : __init__.py
+# @Time         : 2024/6/7 16:58
+# @Author       : betterme
+# @WeChat       : meutils
+# @Software     : PyCharm
+# @Description  : 
+
+from meutils.pipe import *

meutils/fastapi_utils/exceptions/http_error.py

@@ -0,0 +1,72 @@
+import json
+import traceback
+
+from fastapi import status
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from fastapi.exceptions import RequestValidationError, HTTPException
+
+from meutils.notice.feishu import send_message
+
+UNAUTHORIZED_HTTPException = HTTPException(
+    status_code=status.HTTP_401_UNAUTHORIZED,
+    detail={
+        "error": {
+            "message": "",
+            "type": "invalid_request_error",
+            "param": None,
+            "code": "invalid_api_key",
+        }
+    })
+
+
+async def http_error_handler(_: Request, exc: HTTPException) -> JSONResponse:
+    # print(exc)
+    content = {
+        "error":
+            {
+                "message": f"{exc.detail}",
+                "type": "http-error",
+            }
+    }
+    return JSONResponse(
+        content=content,
+        status_code=exc.status_code
+    )
+
+
+async def general_exception_handler(request: Request, exc: Exception):
+    return JSONResponse(
+        content={"message": str(exc)},
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+    )
+
+
+async def chatfire_api_exception_handler(request: Request, exc: Exception):
+    content = {
+        "error":
+            {
+                "message": f"{exc}",
+                "type": "cf-api-error",
+            },
+
+        # "code": status.HTTP_500_INTERNAL_SERVER_ERROR
+    }
+    content_detail = f"{traceback.format_exc()}"
+
+    send_message(
+        content_detail,
+        title="cf-api-error",
+        url="https://open.feishu.cn/open-apis/bot/v2/hook/79fc258f-46a9-419e-b131-1d79b3d0bcff"
+    )
+
+    # send_message
+    return JSONResponse(
+        content=content,
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+    )
+
+
+if __name__ == '__main__':
+    from meutils.notice.feishu import send_message

meutils/fastapi_utils/exceptions/validation_error.py

@@ -0,0 +1,44 @@
+from typing import Union
+
+from fastapi import status
+from fastapi.exceptions import RequestValidationError
+from fastapi.openapi.constants import REF_PREFIX
+from fastapi.openapi.utils import validation_error_response_definition
+from pydantic import ValidationError
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.status import HTTP_422_UNPROCESSABLE_ENTITY
+
+
+async def http422_error_handler(
+        _: Request,
+        exc: Union[RequestValidationError, ValidationError],
+) -> JSONResponse:
+    return JSONResponse(
+        {"errors": exc.errors()},
+        status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+    )
+
+
+validation_error_response_definition["properties"] = {
+    "errors": {
+        "title": "Errors",
+        "type": "array",
+        "items": {"$ref": "{0}ValidationError".format(REF_PREFIX)},
+    },
+}
+
+
+async def validation_exception_handler(_: Request, exc: RequestValidationError) -> JSONResponse:
+    # print(exc)
+    content = {
+        "error":
+            {
+                "message": f"{exc}",
+                "type": "RequestValidationError",
+            }
+    }
+    return JSONResponse(
+        content=content,
+        status_code=status.HTTP_400_BAD_REQUEST
+    )