pykitool 0.0.1__py3-none-any.whl

pykitool/base/tlog.py

@@ -0,0 +1,230 @@
+import builtins
+import inspect
+import os
+import sys
+import warnings
+
+import urllib3
+
+# 忽略
+warnings.filterwarnings("ignore")
+
+# 禁用在使用不安全的 HTTPS 请求时产生的警告
+urllib3.disable_warnings()
+
+# 关闭
+import logging
+import threading
+import time
+
+from loguru import logger
+
+# 格式化
+format = "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level:<5}</level> {thread:<5} <cyan>{file}:{line}</cyan> - {message}"
+# format = "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level:<7}</level> | <yellow>{thread}</yellow> | <cyan>{file}:{function}:{line}</cyan> - <level>{message}</level>"
+
+# 白名单列表
+WHITELIST = []
+
+# ================================ 目录 ================================
+
+# 可替换的日志目录获取函数（默认返回 logs，支持外部注入）
+_log_dir_getter: callable = lambda: "logs"
+
+
+# 默认目录
+def get_log_dir() -> str:
+    return _log_dir_getter()
+
+
+# 供外部注入自定义日志目录获取函数
+def set_log_dir_getter(fn: callable):
+    """注入自定义日志目录获取函数，注入后调用 update_log_level() 立即生效。
+
+    示例::
+
+        # 从配置动态读取
+        set_log_dir_getter(lambda: const.DIR_LOGS)
+        update_log_level()
+
+        # 直接指定固定目录
+        set_log_dir_getter(lambda: "/var/log/myapp")
+        update_log_level()
+    """
+    global _log_dir_getter
+    _log_dir_getter = fn
+
+
+# ================================ 级别 ================================
+
+# 可替换的日志级别获取函数（默认返回 INFO，支持外部注入）
+_log_level_getter: callable = lambda: "INFO"
+
+
+# 默认级别
+def get_log_level() -> str:
+    return _log_level_getter()
+
+
+# 供外部注入自定义日志级别获取函数
+def set_log_level_getter(fn: callable):
+    """注入自定义日志级别获取函数，注入后调用 update_log_level() 立即生效。
+
+    示例::
+
+        # 从配置动态读取
+        set_log_level_getter(lambda: config.LOG_LEVEL)
+        update_log_level()
+
+        # 直接指定固定级别
+        set_log_level_getter(lambda: "DEBUG")
+        update_log_level()
+    """
+    global _log_level_getter
+    _log_level_getter = fn
+
+
+# 拦截器
+class InterceptHandler(logging.Handler):
+    def emit(self, record):
+        # 只处理属于白名单的模块
+        if not any(record.name.startswith(whitelisted) for whitelisted in WHITELIST):
+            if get_log_level() == "DEBUG":
+                print(f"{record.name} -> {record.getMessage()}")
+            return
+        # 尝试获取 loguru 的级别名，失败则使用原始 levelno
+        try:
+            level = logger.level(record.levelname).name
+        except ValueError:
+            level = record.levelno
+        logger.opt(depth=6, exception=record.exc_info).log(level, record.getMessage())
+
+
+# 启动器
+class StartupTimer:
+    def __init__(self):
+        self.start_time = time.perf_counter()
+        self.last_mark = self.start_time
+        self.records = []
+        self.lock = threading.Lock()
+
+    def _get_caller_info(self):
+        frame = inspect.currentframe().f_back.f_back
+        filename = frame.f_code.co_filename
+        lineno = frame.f_lineno
+        return f"{filename}:{lineno}"
+
+    # 跟踪
+    def track(self, tag=None, print: bool = True):
+        with self.lock:
+            now = time.perf_counter()
+            delta = now - self.last_mark
+            total = now - self.start_time
+
+            if tag is None:
+                tag = self._get_caller_info()
+
+            self.records.append((tag, delta, total))
+
+            self.last_mark = now
+            if print:
+                logger.info(f"[StartupTimer] {tag}: +{delta*1000:07.2f} ms (total {total*1000:07.2f} ms)")
+
+    # 计算总时长
+    def total_time(self):
+        return time.perf_counter() - self.start_time
+
+    # 摘要输出
+    def dump(self):
+        for tag, delta, total in self.records:
+            print(f"{tag:12} | +{delta*1000:07.2f} ms | total {total*1000:07.2f} ms")
+
+
+# 初始化
+class LoguruLogger:
+
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            # 单利模式
+            cls._instance = super(LoguruLogger, cls).__new__(cls)
+            cls._instance._configure_logger()
+            # 清空默认自带的日志处理器
+            loggers = [logging.getLogger(name) for name in logging.root.manager.loggerDict]
+            for log in loggers:
+                log.handlers = []
+            # loguru 接管 logging
+            logging.basicConfig(handlers=[InterceptHandler()], level=0)
+            # loguru 接管 print()
+            # builtins.print = logger.info
+            # 设置 comtypes 包的日志级别为 WARNING，屏蔽其 INFO 级别的日志
+        return cls._instance
+
+    def _configure_logger(self):
+        # 获取日志级别
+        current_level = get_log_level()
+        # 获取日志
+        self.logger = logger
+        # 移除默认日志输出
+        self.logger.remove()
+        # 配置控制台输出格式
+        self.logger.add(sys.stdout, level=current_level, format=format, colorize=True)
+        # 配置文件输出并启用轮转 - 按日期轮转（每天午夜）和文件大小轮转
+        log_dir = get_log_dir()
+        os.makedirs(log_dir, exist_ok=True)
+        self.logger.add(log_dir + "/creator_{time:YYYY-MM-DD}.log", rotation="00:00", level=current_level, format=format, colorize=False)
+
+    def update_log_level(self):
+        # 更新日志级别配置
+        self._configure_logger()
+
+
+# 重试
+def retry_print(retry_state):
+    retry = retry_state.retry_object
+    stop = retry.stop
+    max_attempts = getattr(stop, "max_attempt_number", "?")
+    logger.warning(f"Retrying {retry_state.attempt_number}/{max_attempts} time(s), {retry_state.fn.__name__} -> {retry_state.outcome.exception()}, waiting {retry_state.next_action.sleep} seconds before retrying.")
+
+
+# 禁用
+def disable_print():
+
+    # 空方法，以便于禁止输出
+    def no_print(*args, **kwargs):
+        pass
+
+    original_print = builtins.print
+    builtins.print = no_print
+    return original_print
+
+
+# 恢复
+def restore_print(original_print):
+    if original_print:
+        builtins.print = original_print
+
+
+# 更新配置
+def update_cfg():
+    if LoguruLogger._instance:
+        LoguruLogger._instance.update_log_level()
+
+
+if not LoguruLogger._instance:
+    LoguruLogger()
+
+
+if __name__ == "__main__":
+    print("message.{}", 11)
+    logging.info("message %s", 22)
+    # 通过日志输出一些消息
+    logger.trace("message.{}", 0)
+    logger.debug("message.{}", 1)
+    logger.info("message.{}", 2)
+    logger.success("message.{}", 3)
+    logger.warning("message.{}", 4)
+    logger.error("message.{}", 5)
+    logger.critical("message.{}", 6)
+    logger.exception("message.{}", 7)

pykitool/sqliter/exception.py

@@ -0,0 +1,30 @@
+class MissingSessionError(Exception):
+    """
+    Excetion raised for when the user tries to access a database session before
+    it is created.
+    """
+
+    def __init__(self):
+        msg = """
+        No session found! Either you are not currently in a request context,
+        or you need to manually create a session context by using a `db`
+        instance as a context manager e.g.:
+
+        with db():
+            db.session.query(User).all()
+        """
+        super().__init__(msg)
+
+
+class SessionNotInitialisedError(Exception):
+    """
+    Exception raised when the user creates a new DB session without first
+    initialising it.
+    """
+
+    def __init__(self):
+        msg = """
+        Session not initialised! Ensure that DBSessionMiddleware has been 
+        initialised before attempting database access.
+        """
+        super().__init__(msg)

pykitool/sqliter/middleware.py

@@ -0,0 +1,84 @@
+from contextvars import ContextVar
+from typing import Dict, Optional, Union
+
+from sqlalchemy.engine import Engine
+from sqlalchemy.engine.url import URL
+from sqlalchemy.orm import sessionmaker as sessionmaker_
+from sqlmodel import Session, create_engine
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.requests import Request
+from starlette.types import ASGIApp
+
+from pykitool.sqliter.exception import MissingSessionError, SessionNotInitialisedError
+
+
+class sessionmaker(sessionmaker_):
+    def __init__(self, *args, **kwargs):
+        if "class_" not in kwargs:
+            kwargs["class_"] = Session
+        super().__init__(*args, **kwargs)
+
+
+_Session: sessionmaker = None
+_session: ContextVar[Optional[Session]] = ContextVar("_session", default=None)
+
+
+class DBSessionMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: ASGIApp, db_url: Optional[Union[str, URL]] = None, custom_engine: Optional[Engine] = None, engine_args: Dict = None, session_args: Dict = None, commit_on_exit: bool = False):
+        super().__init__(app)
+        global _Session
+        engine_args = engine_args or {}
+        self.commit_on_exit = commit_on_exit
+        session_args = session_args or {}
+        if not custom_engine and not db_url:
+            raise ValueError("You need to pass a db_url or a custom_engine parameter.")
+        if not custom_engine:
+            engine = create_engine(db_url, **engine_args)
+        else:
+            engine = custom_engine
+        _Session = sessionmaker(bind=engine, **session_args)
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
+        with db(commit_on_exit=self.commit_on_exit):
+            response = await call_next(request)
+        return response
+
+
+class DBSessionMeta(type):
+    # using this metaclass means that we can access db.session as a property
+    # at a class level,
+    # rather than db().session
+
+    @property
+    def session(self) -> Session:
+        if _Session is None:
+            raise SessionNotInitialisedError
+        session = _session.get()
+        if session is None:
+            raise MissingSessionError
+        return session
+
+
+class DBSession(metaclass=DBSessionMeta):
+    def __init__(self, session_args: Dict = None, commit_on_exit: bool = False):
+        self.token = None
+        self.session_args = session_args or {}
+        self.commit_on_exit = commit_on_exit
+
+    def __enter__(self):
+        if not isinstance(_Session, sessionmaker):
+            raise SessionNotInitialisedError
+        self.token = _session.set(_Session(**self.session_args))
+        return type(self)
+
+    def __exit__(self, exc_type, *_):
+        sess = _session.get()
+        if exc_type is not None:
+            sess.rollback()
+        if self.commit_on_exit:
+            sess.commit()
+        sess.close()
+        _session.reset(self.token)
+
+
+db: DBSessionMeta = DBSession

pykitool/sqliter/plus.py

@@ -0,0 +1,125 @@
+from typing import Any, Dict, Tuple, Union
+
+import sqlmodel
+from pydantic import ConfigDict
+from sqlalchemy.engine import Engine
+from sqlalchemy.sql import Select
+from sqlmodel import SQLModel, select, text
+from sqlmodel.sql.expression import SelectOfScalar
+
+
+class EngineException(Exception):
+    pass
+
+
+# 类属性描述符：兼容 Python 3.13 移除 @classmethod + @property 链式用法
+class _ClassPropertyDescriptor:
+    def __init__(self, func):
+        self.func = func
+
+    def __get__(self, obj, cls=None):
+        if cls is None:
+            cls = type(obj)
+        return self.func(cls)
+
+
+def classproperty(func):
+    return _ClassPropertyDescriptor(func)
+
+
+class SQLModelPlus(SQLModel):
+    model_config = ConfigDict(ignored_types=(_ClassPropertyDescriptor,))
+
+    __engines__: Dict[str, Engine] = {}
+
+    @classproperty
+    def __get_scope(cls) -> str:
+        return str(cls.__scope__) if hasattr(cls, "__scope__") else "default"
+
+    @classmethod
+    def create_tables(cls, *args, **kwargs):
+        cls.metadata.create_all(cls.__engines__.get(cls.__get_scope), *args, **kwargs)
+
+    @classmethod
+    def find_by_id(cls, ident: Union[Dict[str, Any], Tuple[Any], Any]):
+        with cls.Session() as session:
+            return session.get(cls, ident)
+
+    def save(self):
+        try:
+            return self.create()
+        except:
+            return self.update()
+
+    def create(self):
+        with self.__class__.Session() as session:
+            session.add(self)
+            session.commit()
+            session.refresh(self)
+        return self
+
+    def update(self):
+        with self.__class__.Session() as session:
+            updated_instance = session.merge(self)
+            session.commit()
+            session.refresh(updated_instance)
+        return updated_instance
+
+    def delete(self):
+        with self.__class__.Session() as session:
+            session.delete(self)
+            session.commit()
+        return self
+
+    @classmethod
+    def set_engine(cls, engine: Engine) -> None:
+        cls.__engines__[cls.__get_scope] = engine
+
+    @classmethod
+    def get_engine(cls) -> Engine:
+        return cls.__engines__.get(cls.__get_scope)
+
+    @classmethod
+    def query(
+        cls,
+        statement: Union[SelectOfScalar, str],
+        params: Union[Dict[str, Any], Tuple[Any]] = {},
+    ):
+        return Query(model_cls=cls, statement=statement, params=params)
+
+    # 获取数据库 Session 实例
+    @classmethod
+    def Session(cls) -> sqlmodel.Session:
+        engine: Engine | None = cls.__engines__.get(cls.__get_scope)
+        if engine is None:
+            raise EngineException("Engine is not initialized. Use `.set_engine` method to set engine.")
+        return sqlmodel.Session(bind=engine)
+
+    @classproperty
+    def select(cls) -> SelectOfScalar:
+        return select(cls)
+
+
+class Query:
+    def __init__(
+        self,
+        model_cls: SQLModelPlus,
+        statement: Union[SelectOfScalar, str],
+        params: Union[Dict[str, Any], Tuple[Any], None] = None,
+    ):
+        self.model_cls = model_cls
+        self.statement = statement if isinstance(statement, (SelectOfScalar, Select)) else text(statement)
+        self.params = params
+
+    @property
+    def all(self):
+        with self.model_cls.Session() as session:
+            return session.exec(self.statement, params=self.params).all()
+
+    @property
+    def first(self):
+        with self.model_cls.Session() as session:
+            return session.exec(self.statement, params=self.params).first()
+
+
+__all__ = ["SQLModelPlus", "EngineException"]

pykitool/sqliter/repo.py

@@ -0,0 +1,188 @@
+from contextlib import contextmanager
+from typing import Optional
+
+from sqlmodel import Session, SQLModel, delete, func, select, update
+
+
+@contextmanager
+def reuse_session_or_new(db_engine=None, session: Optional[Session] = None):
+    """
+    Context manager to wrap session reuse or creation logic.
+
+    :param session: An existing session to reuse. If None,
+        a new session is created.
+    :param db_engine: The database engine to use if creating a new session.
+    """
+    should_close = False
+    try:
+        # If session is None, create a new session using the provided db_engine
+        if session is None:
+            if db_engine is None:
+                raise ValueError("No session and no db_engine provided to create a session.")
+            session = Session(db_engine)
+            should_close = True
+
+        # Yield the session for use in the context block
+        yield session
+    finally:
+        # Close the session if it was created inside this context manager
+        if should_close:
+            session.close()
+
+
+class SQLModelRepo:
+    def __init__(self, model: SQLModel, db_engine, init_stmt=None, session=None):
+        """Generic repository for SQLModel.
+
+        Args:
+            model (SQLModel): The SQLModel class (table) for which
+                the repo is instantiated.
+            db_engine: The SQLAlchemy engine linked to the database.
+
+        Usage:
+            users_repo = SQLModelRepo(model=User, db_engine=engine)
+            users_repo.get_by_id(1)
+        """
+        self.model = model
+        self._init_stmt = init_stmt
+        self.db_engine = db_engine
+        self.session = session
+
+    def __call__(self, session):
+        new_repo = SQLModelRepo(model=self.model, db_engine=self.db_engine)
+        new_repo.session = session
+        return new_repo
+
+    def create(self, **kwargs):
+        """Create a new record and save to the database."""
+        instance = self.model(**kwargs)
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            session.add(instance)
+            session.commit()
+            session.refresh(instance)
+        return instance
+
+    def get_by_id(self, id, *fields):
+        """Fetch an object by its primary key."""
+        stmt = self.init_stmt(*fields)
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            return session.exec(stmt.where(getattr(self.model, "id") == id)).first()
+
+    def save(self, instance):
+        """Save the current object (instance) to the database."""
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            session.add(instance)
+            session.commit()
+            session.refresh(instance)
+
+    def save_or_update(self, instance):
+        """Save the current object (instance) to the database."""
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            existing_obj = session.exec(select(self.model).where(self.model.id == instance.id)).first()
+            if existing_obj:
+                for k, v in instance.model_dump().items():
+                    setattr(existing_obj, k, v)
+                    session.add(existing_obj)
+                    instance = existing_obj
+            else:
+                session.add(instance)
+            session.commit()
+            session.refresh(instance)
+
+    def update(self, id, **kwargs):
+        """Record partial update."""
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            update_stmt = update(self.model).where(self.model.id == id).values(**kwargs)
+            session.execute(update_stmt)
+            session.commit()
+
+    def update_all(self, **kwargs):
+        """Partial update for all selected records."""
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            if self._init_stmt:
+                update_stmt = update(self.model).where(self.init_stmt().whereclause).values(**kwargs)
+            else:
+                update_stmt = update(self.model).values(**kwargs)
+            session.execute(update_stmt)
+            session.commit()
+
+    def delete(self, instance):
+        """Delete an object from the database."""
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            session.delete(instance)
+            session.commit()
+
+    def delete_all(self):
+        """Delete all records in query."""
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            if self._init_stmt:
+                delete_stmt = delete(self.model).where(self.init_stmt().whereclause)
+            else:
+                delete_stmt = delete(self.model)
+            session.execute(delete_stmt)
+            session.commit()
+
+    def filter(self, *filters, _fields=(), **kwargs) -> "SQLModelRepo":
+        """Filter records based on provided conditions."""
+        stmt = self.init_stmt(*_fields).where(*filters, *[getattr(self.model, k) == v if isinstance(k, str) else k == v for k, v in kwargs.items()])
+        return SQLModelRepo(init_stmt=stmt, model=self.model, db_engine=self.db_engine, session=self.session)
+
+    def paginate(self, offset: int, limit: int, order_by: str, desc: bool = False) -> list:
+        """Paginate results"""
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            return self._paginate(session, offset, limit, order_by, desc)
+
+    def paginate_with_total(self, offset: int, limit: int, order_by: str, desc: bool = False) -> (list, int):  # type: ignore
+        """Paginate results and fetch total count
+
+        Returns:
+            tuple(list, int) - Items and total count.
+        """
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            count_stmt = select(func.count()).select_from(self.init_stmt().subquery())
+            count = session.execute(count_stmt).scalar()
+            results = self._paginate(session, offset, limit, order_by, desc)
+            return results, count
+
+    def _paginate(self, session, offset: int, limit: int, order_by: str, desc: bool = False) -> list:
+        order_by = getattr(self.model, order_by)
+        if desc:
+            order_by = getattr(order_by, "desc")()
+        return session.exec(self.init_stmt().order_by(order_by).offset(offset).limit(limit)).all()
+
+    def all(self) -> list:
+        """Get all results"""
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            return session.exec(self.init_stmt()).all()
+
+    def count(self):
+        """Get total results count"""
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            count_stmt = select(func.count()).select_from(self.init_stmt().subquery())
+            return session.execute(count_stmt).scalar()
+
+    def first(self):
+        with reuse_session_or_new(self.db_engine, self.session) as session:
+            return session.exec(self.init_stmt()).first()
+
+    def get_or_404(self, id):
+        if not (obj := self.get_by_id(id)):
+            raise Exception(status_code=404, detail=f"{self.model.__name__.title()} with id {id} not found")
+        return obj
+
+    def delete_or_404(self, id):
+        obj = self.get_or_404(id)
+        self.delete(obj)
+
+    def update_or_404(self, id, **kwargs):
+        if self.get_or_404(id):
+            self.update(id, **kwargs)
+
+    def _get_select_obj(self, fields=None):
+        return [self.model] if not fields else [getattr(self.model, f) for f in fields]
+
+    def init_stmt(self, *fields):
+        if self._init_stmt is not None:
+            return self._init_stmt
+        else:
+            return select(*self._get_select_obj(fields))