python-basekit 0.0.11__py3-none-any.whl

basekit/cache/schema.py

@@ -0,0 +1,53 @@
+from collections.abc import Awaitable, Callable
+from contextlib import AbstractAsyncContextManager, suppress
+from functools import wraps
+
+from pydantic import BaseModel
+
+from basekit.cache.utils import (
+    CacheValueError,
+    deserialize_cache_value,
+    get_instance,
+    serialize_cache_key,
+    serialize_cache_value,
+)
+
+
+class Cache(BaseModel):
+    def lifespan(self) -> AbstractAsyncContextManager[None]:
+        raise NotImplementedError
+
+    async def set(self, key: bytes, value: bytes, /) -> None:
+        raise NotImplementedError
+
+    async def get(self, key: bytes, /) -> bytes | None:
+        raise NotImplementedError
+
+    async def delete(self, key: bytes, /) -> None:
+        raise NotImplementedError
+
+    def __call__[**P, T](
+        self, func: Callable[P, Awaitable[T]], /
+    ) -> Callable[P, Awaitable[T]]:
+        @wraps(wrapped=func)
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            data: dict = {
+                "module": func.__module__,
+                "function": func.__qualname__,
+                "instance": get_instance(func),
+                "args": args,
+                "kwargs": kwargs,
+            }
+            key: bytes = serialize_cache_key(data)
+
+            cached_value: bytes | None = await self.get(key)
+            if cached_value is not None:
+                with suppress(CacheValueError):
+                    return deserialize_cache_value(cached_value)
+                await self.delete(key)
+
+            value: T = await func(*args, **kwargs)
+            await self.set(key, serialize_cache_value(value))
+            return value
+
+        return wrapper

basekit/cache/utils.py

@@ -0,0 +1,113 @@
+import pickle
+import pickletools
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class CacheKeyError(Exception):
+    pass
+
+
+class CacheValueError(Exception):
+    pass
+
+
+def get_instance(func: Callable, /) -> Any:
+    current: Callable | None = func
+    while current is not None:
+        instance: Any = getattr(current, "__self__", None)
+        if instance is not None:
+            return instance
+        current = getattr(current, "__wrapped__", None)
+    return None
+
+
+def cache_key_fields(cls: type[BaseModel], /) -> list[str]:
+    model_fields: set[str] = set()
+    cache_key_fields: set[str] = set()
+    for base in reversed(cls.__mro__):
+        if not issubclass(base, BaseModel):
+            continue
+        current_model_fields: set[str] = set(base.model_fields.keys()) - model_fields
+        current_cache_key_fields: set[str] = set(
+            base.__dict__.get("cache_key_fields", current_model_fields)
+        )
+        model_fields.update(current_model_fields)
+        cache_key_fields.update(current_cache_key_fields)
+    return list(cache_key_fields)
+
+
+def build_cache_key(value: Any, /) -> dict:
+    if isinstance(value, type) and issubclass(value, BaseModel):
+        return {
+            "module": value.__module__,
+            "class": value.__qualname__,
+            "schema": value.model_json_schema(),
+        }
+    if isinstance(value, BaseModel):
+        cls: type[BaseModel] = value.__class__
+        instance: dict = {f: getattr(value, f) for f in cache_key_fields(cls)}
+        return {
+            "module": cls.__module__,
+            "class": cls.__qualname__,
+            "instance": build_cache_key(instance),
+        }
+    if isinstance(value, dict):
+        if not all(isinstance(k, str) for k in value.keys()):
+            raise ValueError(f"字典键必须是字符串：{value.keys()}")
+        return {"dict": {k: build_cache_key(v) for k, v in value.items()}}
+    if isinstance(value, list):
+        return {"list": [build_cache_key(i) for i in value]}
+    if isinstance(value, tuple):
+        return {"tuple": [build_cache_key(i) for i in value]}
+    if isinstance(value, bytes):
+        return {"bytes": value}
+    if isinstance(value, str):
+        return {"str": value}
+    if isinstance(value, bool):
+        return {"bool": value}
+    if isinstance(value, int):
+        return {"int": value}
+    if isinstance(value, float):
+        return {"float": value}
+    if value is None:
+        return {"none": None}
+    raise ValueError(f"类型不支持序列化：{value.__class__}")
+
+
+def normalize_cache_key(value: Any, /) -> Any:
+    if isinstance(value, dict):
+        return {k: normalize_cache_key(value[k]) for k in sorted(value.keys())}
+    if isinstance(value, list):
+        return [normalize_cache_key(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(normalize_cache_key(item) for item in value)
+    return value
+
+
+def serialize_cache_key(value: Any, /) -> bytes:
+    try:
+        data: dict = normalize_cache_key(build_cache_key(value))
+        return pickletools.optimize(
+            p=pickle.dumps(obj=data, protocol=pickle.HIGHEST_PROTOCOL)
+        )
+    except Exception as exc:
+        raise CacheKeyError("缓存键序列化失败") from exc
+
+
+def serialize_cache_value(value: Any, /) -> bytes:
+    try:
+        return pickletools.optimize(
+            p=pickle.dumps(obj=value, protocol=pickle.HIGHEST_PROTOCOL)
+        )
+    except Exception as exc:
+        raise CacheValueError("缓存值序列化失败") from exc
+
+
+def deserialize_cache_value(value: bytes, /) -> Any:
+    try:
+        return pickle.loads(value)
+    except Exception as exc:
+        raise CacheValueError("缓存值反序列化失败") from exc

basekit/database.py

@@ -0,0 +1,33 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import override
+
+from pydantic import BaseModel, PrivateAttr
+from sqlalchemy import MetaData
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
+
+
+class DatabaseClient(BaseModel):
+    url: str
+
+    _engine: AsyncEngine = PrivateAttr()
+
+    @override
+    def model_post_init(self, context) -> None:
+        self._engine = create_async_engine(url=self.url)
+
+    @asynccontextmanager
+    async def lifespan(self) -> AsyncIterator[None]:
+        try:
+            yield
+        finally:
+            await self._engine.dispose()
+
+    async def create_schema(self, metadata: MetaData, /) -> None:
+        async with self._engine.begin() as conn:
+            await conn.run_sync(fn=metadata.create_all)
+
+    @asynccontextmanager
+    async def connection(self) -> AsyncIterator[AsyncConnection]:
+        async with self._engine.connect() as connection:
+            yield connection

basekit/http/clients/curl_cffi.py

@@ -0,0 +1,99 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager, suppress
+from http.cookiejar import CookieJar
+from typing import Any, override
+
+import logfire
+from curl_cffi import AsyncSession, BrowserTypeLiteral, Response
+from curl_cffi.requests.exceptions import HTTPError
+from pydantic import PrivateAttr
+
+from basekit.http.schema import HttpClient, HttpMethod
+from basekit.http.utils import clone_cookiejar, truncate_data
+
+
+class CurlCffiClient(HttpClient[Response]):
+    impersonate: BrowserTypeLiteral = "chrome"
+
+    _client: AsyncSession = PrivateAttr()
+
+    @override
+    def model_post_init(self, context: Any) -> None:
+        self._client = AsyncSession(
+            headers=self.headers,
+            verify=self.verify,
+            http_version="v2",
+            timeout=self.timeout,
+            allow_redirects=True,
+            max_clients=self.limit,
+            impersonate=self.impersonate,
+        )
+
+    @override
+    @asynccontextmanager
+    async def lifespan(self) -> AsyncIterator[None]:
+        async with self._client:
+            yield
+
+    def save_cookies(self) -> CookieJar:
+        return clone_cookiejar(self._client.cookies.jar)
+
+    def load_cookies(self, cookies: CookieJar) -> None:
+        self._client.cookies = clone_cookiejar(cookies)
+
+    def _extract_content(self, response: Response, /) -> dict | str | bytes:
+        content_type: str = response.headers.get(key="Content-Type", default="") or ""
+        content_type = content_type.split(sep=";")[0].strip().lower()
+        if content_type == "application/json":
+            with suppress(Exception):
+                return response.json()
+        if content_type.startswith("text/"):
+            with suppress(Exception):
+                return response.text
+        with suppress(Exception):
+            return response.content
+        return b""
+
+    async def _request_once(
+        self,
+        method: HttpMethod,
+        url: str,
+        /,
+        *,
+        params: dict | None = None,
+        data: dict | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> Response:
+        with logfire.span(f"curl cffi client | request | {method} | {url}") as span:
+            span.set_attribute(key="request_params", value=truncate_data(params))
+            span.set_attribute(key="request_data", value=truncate_data(data))
+            span.set_attribute(key="request_headers", value=truncate_data(headers))
+
+            response: Response = await self._client.request(
+                method=method, url=url, params=params, json=data, headers=headers
+            )
+
+            status_code: int = response.status_code
+            span.set_attribute(key="status_code", value=status_code)
+            span.message = f"{span.message} -> {status_code}"
+
+            response_headers: dict = dict(response.headers.items())
+            span.set_attribute(
+                key="response_headers", value=truncate_data(response_headers)
+            )
+
+            response_content: dict | str | bytes = self._extract_content(response)
+            span.set_attribute(
+                key="response_content", value=truncate_data(response_content)
+            )
+
+            response.raise_for_status()
+            return response
+
+    def _should_retry(self, exc: BaseException, /) -> bool:
+        if not isinstance(exc, HTTPError):
+            return True
+        if exc.response is None:
+            return True
+        status_code: int = exc.response.status_code
+        return status_code == 429 or status_code >= 500

basekit/http/clients/httpx.py

@@ -0,0 +1,100 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager, suppress
+from http.cookiejar import CookieJar
+from typing import Any, override
+
+import logfire
+from httpx import AsyncClient, HTTPStatusError, Limits, Response
+from pydantic import PrivateAttr
+
+from basekit.http.schema import HttpClient, HttpMethod
+from basekit.http.utils import clone_cookiejar, truncate_data
+
+
+class HttpxClient(HttpClient[Response]):
+    _client: AsyncClient = PrivateAttr()
+
+    @override
+    def model_post_init(self, context: Any) -> None:
+        self._client = AsyncClient(
+            headers=self.headers,
+            verify=self.verify,
+            http2=True,
+            timeout=self.timeout,
+            follow_redirects=True,
+            limits=Limits(
+                max_connections=self.limit,
+                max_keepalive_connections=self.limit,
+            ),
+        )
+
+    @override
+    @asynccontextmanager
+    async def lifespan(self) -> AsyncIterator[None]:
+        async with self._client:
+            yield
+
+    @override
+    def save_cookies(self) -> CookieJar:
+        return clone_cookiejar(self._client.cookies.jar)
+
+    @override
+    def load_cookies(self, cookies: CookieJar) -> None:
+        self._client.cookies = clone_cookiejar(cookies)
+
+    def _extract_content(self, response: Response, /) -> dict | str | bytes:
+        content_type: str = response.headers.get(key="Content-Type", default="") or ""
+        content_type = content_type.split(sep=";")[0].strip().lower()
+        if content_type == "application/json":
+            with suppress(Exception):
+                return response.json()
+        if content_type.startswith("text/"):
+            with suppress(Exception):
+                return response.text
+        with suppress(Exception):
+            return response.content
+        return b""
+
+    @override
+    async def _request_once(
+        self,
+        method: HttpMethod,
+        url: str,
+        /,
+        *,
+        params: dict | None = None,
+        data: dict | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> Response:
+        with logfire.span(f"httpx client | request | {method} | {url}") as span:
+            span.set_attribute(key="request_params", value=truncate_data(params))
+            span.set_attribute(key="request_data", value=truncate_data(data))
+            span.set_attribute(key="request_headers", value=truncate_data(headers))
+
+            response: Response = await self._client.request(
+                method=method, url=url, params=params, json=data, headers=headers
+            )
+
+            status_code: int = response.status_code
+            span.set_attribute(key="status_code", value=status_code)
+            span.message = f"{span.message} -> {status_code}"
+
+            response_headers: dict = dict(response.headers.items())
+            span.set_attribute(
+                key="response_headers", value=truncate_data(response_headers)
+            )
+
+            response_content: dict | str | bytes = self._extract_content(response)
+            span.set_attribute(
+                key="response_content", value=truncate_data(response_content)
+            )
+
+            response.raise_for_status()
+            return response
+
+    @override
+    def _should_retry(self, exc: BaseException, /) -> bool:
+        if not isinstance(exc, HTTPStatusError):
+            return True
+        status_code: int = exc.response.status_code
+        return status_code == 429 or status_code >= 500

basekit/http/schema.py

@@ -0,0 +1,111 @@
+from contextlib import AbstractAsyncContextManager
+from http.cookiejar import CookieJar
+from typing import ClassVar, Literal
+
+from pydantic import BaseModel, ConfigDict
+from tenacity import (
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential_jitter,
+)
+from tenacity.stop import stop_base
+from tenacity.wait import wait_base
+
+DEFAULT_TIMEOUT = 5.0
+DEFAULT_LIMIT = 100
+DEFAULT_RETRY_STOP = stop_after_attempt(max_attempt_number=3)
+DEFAULT_RETRY_WAIT = wait_exponential_jitter()
+
+type HttpMethod = Literal["GET", "POST"]
+
+
+class HttpClient[ResponseT](BaseModel):
+    model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
+
+    headers: dict[str, str] | None = None
+    verify: bool = True
+    timeout: float = DEFAULT_TIMEOUT
+    limit: int = DEFAULT_LIMIT
+    retry_stop: stop_base = DEFAULT_RETRY_STOP
+    retry_wait: wait_base = DEFAULT_RETRY_WAIT
+
+    def lifespan(self) -> AbstractAsyncContextManager[None]:
+        raise NotImplementedError
+
+    def save_cookies(self) -> CookieJar:
+        raise NotImplementedError
+
+    def load_cookies(self, cookies: CookieJar) -> None:
+        raise NotImplementedError
+
+    async def _request_once(
+        self,
+        method: HttpMethod,
+        url: str,
+        /,
+        *,
+        params: dict | None = None,
+        data: dict | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> ResponseT:
+        raise NotImplementedError
+
+    def _should_retry(self, exc: BaseException, /) -> bool:
+        raise NotImplementedError
+
+    async def request(
+        self,
+        method: HttpMethod,
+        url: str,
+        /,
+        *,
+        params: dict | None = None,
+        data: dict | None = None,
+        headers: dict[str, str] | None = None,
+        retry_on_failure: bool = False,
+    ) -> ResponseT:
+        if not retry_on_failure:
+            return await self._request_once(
+                method, url, params=params, data=data, headers=headers
+            )
+        return await retry(
+            retry=retry_if_exception(predicate=self._should_retry),
+            stop=self.retry_stop,
+            wait=self.retry_wait,
+            reraise=True,
+        )(self._request_once)(method, url, params=params, data=data, headers=headers)
+
+    async def get(
+        self,
+        url: str,
+        /,
+        *,
+        params: dict | None = None,
+        headers: dict[str, str] | None = None,
+        retry_on_failure: bool = True,
+    ) -> ResponseT:
+        return await self.request(
+            "GET",
+            url,
+            params=params,
+            headers=headers,
+            retry_on_failure=retry_on_failure,
+        )
+
+    async def post(
+        self,
+        url: str,
+        /,
+        *,
+        data: dict | None = None,
+        headers: dict[str, str] | None = None,
+        retry_on_failure: bool = False,
+    ) -> ResponseT:
+        return await self.request(
+            "POST",
+            url,
+            data=data,
+            headers=headers,
+            retry_on_failure=retry_on_failure,
+        )

basekit/http/utils.py

@@ -0,0 +1,31 @@
+from http.cookiejar import CookieJar
+from typing import Any
+
+SIZE_LIMIT = 1048576  # 1MB
+
+
+def truncate_data(value: Any, /) -> Any:
+    match value:
+        case dict():
+            return {k: truncate_data(v) for k, v in value.items()}
+        case list():
+            return [truncate_data(v) for v in value]
+        case str():
+            size: int = len(value)
+            if size <= SIZE_LIMIT:
+                return value
+            return f"[TOO LARGE: {size} chars]{value[:SIZE_LIMIT]}..."
+        case bytes():
+            size: int = len(value)
+            if size <= SIZE_LIMIT:
+                return value
+            return f"[TOO LARGE: {size} bytes]{value[:SIZE_LIMIT]}..."
+        case _:
+            return value
+
+
+def clone_cookiejar(cookiejar: CookieJar, /) -> CookieJar:
+    cloned_cookie_jar = CookieJar()
+    for cookie in cookiejar:
+        cloned_cookie_jar.set_cookie(cookie=cookie)
+    return cloned_cookie_jar