kotonebot 0.1.0__py3-none-any.whl

kotonebot/backend/context/task_action.py

@@ -0,0 +1,176 @@
+import logging
+from typing import Callable, ParamSpec, TypeVar, overload, Literal
+from dataclasses import dataclass
+
+
+from .context import ContextStackVars, ScreenshotMode
+from ...errors import TaskNotFoundError
+
+P = ParamSpec('P')
+R = TypeVar('R')
+logger = logging.getLogger(__name__)
+
+TaskRunAtType = Literal['pre', 'post', 'manual', 'regular'] | str
+
+
+@dataclass
+class Task:
+    name: str
+    id: str
+    description: str
+    func: Callable
+    priority: int
+    """
+    任务优先级，数字越大优先级越高。
+    """
+    run_at: TaskRunAtType = 'regular'
+
+
+@dataclass
+class Action:
+    name: str
+    description: str
+    func: Callable
+    priority: int
+    """
+    动作优先级，数字越大优先级越高。
+    """
+
+
+task_registry: dict[str, Task] = {}
+action_registry: dict[str, Action] = {}
+current_callstack: list[Task|Action] = []
+
+def _placeholder():
+    raise NotImplementedError('Placeholder function')
+
+def task(
+    name: str,
+    task_id: str|None = None,
+    description: str|None = None,
+    *,
+    pass_through: bool = False,
+    priority: int = 0,
+    screenshot_mode: ScreenshotMode = 'auto',
+    run_at: TaskRunAtType = 'regular'
+):
+    """
+    `task` 装饰器，用于标记一个函数为任务函数。
+
+    :param name: 任务名称
+    :param task_id: 任务 ID。如果为 None，则使用函数名称作为 ID。
+    :param description: 任务描述。如果为 None，则使用函数的 docstring 作为描述。
+    :param pass_through: 
+        默认情况下， @task 装饰器会包裹任务函数，跟踪其执行情况。
+        如果不想跟踪，则设置此参数为 False。
+    :param priority: 任务优先级，数字越大优先级越高。
+    :param run_at: 任务运行时间。
+    """
+    # 设置 ID
+    # 获取 caller 信息
+    def _task_decorator(func: Callable[P, R]) -> Callable[P, R]:
+        nonlocal description, task_id
+        description = description or func.__doc__ or ''
+        # TODO: task_id 冲突检测
+        task_id = task_id or func.__name__
+        task = Task(name, task_id, description, _placeholder, priority, run_at)
+        task_registry[name] = task
+        logger.debug(f'Task "{name}" registered.')
+        if pass_through:
+            return func
+        else:
+            def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+                current_callstack.append(task)
+                vars = ContextStackVars.push(screenshot_mode=screenshot_mode)
+                ret = func(*args, **kwargs)
+                ContextStackVars.pop()
+                current_callstack.pop()
+                return ret
+            task.func = _wrapper
+            return _wrapper
+    return _task_decorator
+
+@overload
+def action(func: Callable[P, R]) -> Callable[P, R]: ...
+
+@overload
+def action(
+    name: str,
+    *,
+    description: str|None = None,
+    pass_through: bool = False,
+    priority: int = 0,
+    screenshot_mode: ScreenshotMode | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """
+    `action` 装饰器，用于标记一个函数为动作函数。
+
+    :param name: 动作名称。如果为 None，则使用函数的名称作为名称。
+    :param description: 动作描述。如果为 None，则使用函数的 docstring 作为描述。
+    :param pass_through: 
+        默认情况下， @action 装饰器会包裹动作函数，跟踪其执行情况。
+        如果不想跟踪，则设置此参数为 False。
+    :param priority: 动作优先级，数字越大优先级越高。
+    :param screenshot_mode: 截图模式。
+    """
+    ...
+
+# TODO: 需要找个地方统一管理这些属性名
+ATTR_ORIGINAL_FUNC = '_kb_inner'
+ATTR_ACTION_MARK = '__kb_action_mark'
+def action(*args, **kwargs):
+    def _register(func: Callable, name: str, description: str|None = None, priority: int = 0) -> Action:
+        description = description or func.__doc__ or ''
+        action = Action(name, description, func, priority)
+        action_registry[name] = action
+        logger.debug(f'Action "{name}" registered.')
+        return action
+
+    if len(args) == 1 and isinstance(args[0], Callable):
+        func = args[0]
+        action = _register(_placeholder, func.__name__, func.__doc__)
+        def _wrapper(*args: P.args, **kwargs: P.kwargs):
+            current_callstack.append(action)
+            vars = ContextStackVars.push()
+            ret = func(*args, **kwargs)
+            ContextStackVars.pop()
+            current_callstack.pop()
+            return ret
+        setattr(_wrapper, ATTR_ORIGINAL_FUNC, func)
+        setattr(_wrapper, ATTR_ACTION_MARK, True)
+        action.func = _wrapper
+        return _wrapper
+    else:
+        name = args[0]
+        description = kwargs.get('description', None)
+        pass_through = kwargs.get('pass_through', False)
+        priority = kwargs.get('priority', 0)
+        screenshot_mode = kwargs.get('screenshot_mode', None)
+        def _action_decorator(func: Callable):
+            nonlocal pass_through
+            action = _register(_placeholder, name, description)
+            pass_through = kwargs.get('pass_through', False)
+            if pass_through:
+                return func
+            else:
+                def _wrapper(*args: P.args, **kwargs: P.kwargs):
+                    current_callstack.append(action)
+                    vars = ContextStackVars.push(screenshot_mode=screenshot_mode)
+                    ret = func(*args, **kwargs)
+                    ContextStackVars.pop()
+                    current_callstack.pop()
+                    return ret
+                setattr(_wrapper, ATTR_ORIGINAL_FUNC, func)
+                setattr(_wrapper, ATTR_ACTION_MARK, True)
+                action.func = _wrapper
+                return _wrapper
+        return _action_decorator
+
+def tasks_from_id(task_ids: list[str]) -> list[Task]:
+    result = []
+    for tid in task_ids:
+        target = next(task for task in task_registry.values() if task.id == tid)
+        if target is None:
+            raise TaskNotFoundError(f'Task "{tid}" not found.')
+        result.append(target)
+    return result

kotonebot/backend/core.py

@@ -0,0 +1,126 @@
+import logging
+from functools import cache
+from typing import Callable
+
+import cv2
+from cv2.typing import MatLike
+
+from kotonebot.util import cv2_imread
+from kotonebot.primitives import RectTuple, Rect, Point
+from kotonebot.errors import ResourceFileMissingError
+
+class Ocr:
+    def __init__(
+        self,
+        text: str | Callable[[str], bool],
+        *,
+        language: str = 'jp',
+    ):
+        self.text = text
+        self.language = language
+
+
+class Image:
+    def __init__(
+        self,
+        *,
+        path: str | None = None,
+        name: str | None = 'untitled',
+        data: MatLike | None = None,
+    ):
+        self.path = path
+        self.name = name
+        self.__data: MatLike | None = data
+        self.__data_with_alpha: MatLike | None = None
+
+    @cache
+    def binary(self) -> 'Image':
+        return Image(data=cv2.cvtColor(self.data, cv2.COLOR_BGR2GRAY))
+
+    @property
+    def data(self) -> MatLike:
+        if self.__data is None:
+            if self.path is None:
+                raise ValueError('Either path or data must be provided.')
+            self.__data = cv2_imread(self.path)
+            if self.__data is None:
+                raise ResourceFileMissingError(self.path, 'sprite')
+            logger.debug(f'Read image "{self.name}" from {self.path}')
+        return self.__data
+    
+    @property
+    def data_with_alpha(self) -> MatLike:
+        if self.__data_with_alpha is None:
+            if self.path is None:
+                raise ValueError('Either path or data must be provided.')
+            self.__data_with_alpha = cv2_imread(self.path, cv2.IMREAD_UNCHANGED)
+            if self.__data_with_alpha is None:
+                raise ResourceFileMissingError(self.path, 'sprite with alpha')
+            logger.debug(f'Read image "{self.name}" from {self.path}')
+        return self.__data_with_alpha
+    
+    def __repr__(self) -> str:
+        if self.path is None:
+            return f'<Image: memory>'
+        else:
+            return f'<Image: "{self.name}" at {self.path}>'
+
+
+class HintBox(Rect):
+    def __init__(
+        self,
+        x1: int,
+        y1: int,
+        x2: int,
+        y2: int,
+        *,
+        name: str | None = None,
+        description: str | None = None,
+        source_resolution: tuple[int, int],
+    ):
+        super().__init__(x1, y1, x2 - x1, y2 - y1, name=name)
+        self.description = description
+        self.source_resolution = source_resolution
+
+    @property
+    def width(self) -> int:
+        return self.x2 - self.x1
+    
+    @property
+    def height(self) -> int:
+        return self.y2 - self.y1
+    
+    @property
+    def rect(self) -> RectTuple:
+        return self.x1, self.y1, self.width, self.height
+
+class HintPoint(Point):
+    def __init__(self, x: int, y: int, *, name: str | None = None, description: str | None = None):
+        super().__init__(x, y, name=name)
+        self.description = description
+
+    def __repr__(self) -> str:
+        return f'HintPoint<"{self.name}" at ({self.x}, {self.y})>'
+
+def unify_image(image: MatLike | str | Image, transparent: bool = False) -> MatLike:
+    if isinstance(image, str):
+        if not transparent:
+            image = cv2_imread(image)
+        else:
+            image = cv2_imread(image, cv2.IMREAD_UNCHANGED)
+    elif isinstance(image, Image):
+        if transparent:
+            image = image.data_with_alpha
+        else:
+            image = image.data
+    return image
+
+logger = logging.getLogger(__name__)
+
+
+if __name__ == '__main__':
+    hint_box = HintBox(100, 100, 200, 200, source_resolution=(1920, 1080))
+    print(hint_box.rect)
+    print(hint_box.width)
+    print(hint_box.height)
+

kotonebot/backend/debug/__init__.py

@@ -0,0 +1 @@
+from .vars import result, debug, img, color

kotonebot/backend/debug/entry.py

@@ -0,0 +1,89 @@
+import os
+import runpy
+import shutil
+import argparse
+import importlib
+from pathlib import Path
+from threading import Thread
+
+from . import debug
+from kotonebot import logging
+from kotonebot.backend.context import init_context
+
+logger = logging.getLogger(__name__)
+
+def _task_thread(task_module: str):
+    """任务线程。"""
+    runpy.run_module(task_module, run_name="__main__")
+
+def _parse_args():
+    """解析命令行参数。"""
+    parser = argparse.ArgumentParser(description='KotoneBot visual debug tool')
+    parser.add_argument(
+        '-s', '--save', 
+        help='Save dump image and results to the specified folder',
+        type=str,
+        metavar='PATH'
+    )
+    parser.add_argument(
+        '-c', '--clear',
+        help='Clear the dump folder before running',
+        action='store_true'
+    )
+    parser.add_argument(
+        '-t', '--config-type',
+        help='The full path of the config data type. e.g. `kotonebot.tasks.common.BaseConfig`',
+        type=str,
+        metavar='TYPE',
+        required=True
+    )
+    parser.add_argument(
+        'input_module',
+        help='The module to run'
+    )
+    return parser.parse_args()
+
+def _start_task_thread(module: str):
+    """启动任务线程。"""
+    thread = Thread(target=_task_thread, args=(module,))
+    thread.start()
+
+if __name__ == "__main__":
+    args = _parse_args()
+    debug.enabled = True
+    
+    # 设置保存路径
+    if args.save:
+        save_path = Path(args.save)
+        debug.auto_save_to_folder = str(save_path)
+        if not os.path.exists(save_path):
+            os.makedirs(save_path)
+    if args.clear:
+        if debug.auto_save_to_folder:
+            try:
+                logger.info(f"Removing {debug.auto_save_to_folder}")
+                shutil.rmtree(debug.auto_save_to_folder)
+            except PermissionError:
+                logger.warning(f"Failed to remove {debug.auto_save_to_folder}. Trying to remove all contents instead.")
+                for root, dirs, files in os.walk(debug.auto_save_to_folder):
+                    for file in files:
+                        try:
+                            os.remove(os.path.join(root, file))
+                        except PermissionError:
+                            raise
+                
+
+    # 初始化上下文
+    module_name, class_name = args.config_type.rsplit('.', 1)
+    class_ = importlib.import_module(module_name).__getattribute__(class_name)
+    init_context(config_type=class_)
+    
+    # 启动服务器
+    from .server import app
+    import uvicorn
+    
+    # 启动任务线程
+    _start_task_thread(args.input_module)
+    
+    # 启动服务器
+    uvicorn.run(app, host="127.0.0.1", port=8000, log_level='critical' if debug.hide_server_log else None)

kotonebot/backend/debug/mock.py

@@ -0,0 +1,79 @@
+import time
+from typing_extensions import override
+
+import cv2
+from cv2.typing import MatLike
+
+from kotonebot import sleep
+from kotonebot.client.device import Device
+
+class Video:
+    def __init__(self, path: str, fps: int):
+        self.path = path
+        self.fps = fps
+        self.paused = False
+        """是否暂停"""
+        self.__cap = cv2.VideoCapture(path)
+        self.__last_frame = None
+        self.__last_time = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.paused:
+            return self.__last_frame
+        ret, frame = self.__cap.read()
+        if not ret:
+            raise StopIteration
+        self.__last_frame = frame
+        self.__last_time = time.time()
+        if self.__last_time - time.time() < 1 / self.fps:
+            sleep(1 / self.fps)
+        return frame
+
+    def pause(self):
+        self.paused = True
+
+    def resume(self):
+        self.paused = False
+
+class MockDevice(Device):
+    def __init__(
+        self
+    ):
+        super().__init__()
+        self.__video_stream = None
+        self.__image = None
+        self.__screen_size = None
+
+    def load_video(self, path: str, fps: int):
+        self.__video_stream = Video(path, fps)
+        return self.__video_stream
+
+    def load_image(self, img: str | MatLike):
+        if isinstance(img, str):
+            self.__image = cv2.imread(img)
+        else:
+            self.__image = img
+        return self.__image
+
+    def set_screen_size(self, width: int, height: int):
+        self.__screen_size = (width, height)
+
+    @override
+    def screenshot(self):
+        if self.__image is not None:
+            return self.__image
+        elif self.__video_stream is not None:
+            return next(self.__video_stream)
+        else:
+            raise RuntimeError('No video stream loaded')
+        
+    @property
+    @override
+    def screen_size(self):
+        if self.__screen_size is not None:
+            return self.__screen_size
+        else:
+            raise RuntimeError('No screen size set')

kotonebot/backend/debug/server.py

@@ -0,0 +1,223 @@
+import time
+import asyncio
+import inspect
+import threading
+import traceback
+import subprocess
+from io import StringIO
+from pathlib import Path
+from typing import Literal
+from collections import deque
+from contextlib import redirect_stdout
+
+import cv2
+import uvicorn
+from thefuzz import fuzz
+from pydantic import BaseModel
+from fastapi.responses import FileResponse, Response
+from fastapi import FastAPI, WebSocket, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+
+import kotonebot
+import kotonebot.backend
+import kotonebot.backend.context
+from kotonebot.backend.core import HintBox, Image
+from ..context import manual_context
+from . import vars as debug_vars
+from .vars import WSImage, WSMessageData, WSMessage, WSCallstack
+
+app = FastAPI()
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
+
+# 获取当前文件夹路径
+CURRENT_DIR = Path(__file__).parent
+
+APP_DIR = Path.cwd()
+
+class File(BaseModel):
+    name: str
+    full_path: str
+    type: Literal["file", "dir"]
+
+@app.get("/api/read_file")
+async def read_file(path: str):
+    """读取文件内容"""
+    try:
+        # 确保路径在当前目录下
+        full_path = (APP_DIR / path).resolve()
+        if not Path(full_path).is_relative_to(APP_DIR):
+            raise HTTPException(status_code=403, detail="Access denied")
+        
+        if not full_path.exists():
+            raise HTTPException(status_code=404, detail="File not found")
+        # 添加缓存控制头
+        headers = {
+            "Cache-Control": "public, max-age=3600",  # 缓存1小时
+            "ETag": f'"{hash(full_path)}"'  # 使用full_path的哈希值作为ETag
+        }
+        return FileResponse(full_path, headers=headers)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/api/read_memory")
+async def read_memory(key: str):
+    """读取内存中的数据"""
+    try:
+        image = None
+        if (image := debug_vars._read_image(key)) is not None:
+            pass
+        else:
+            raise HTTPException(status_code=404, detail="Key not found")
+        
+        # 编码图片
+        encode_params = [cv2.IMWRITE_PNG_COMPRESSION, 4]
+        _, buffer = cv2.imencode('.png', image, encode_params)
+        # 添加缓存控制头
+        headers = {
+            "Cache-Control": "public, max-age=3600",  # 缓存1小时
+            "ETag": f'"{hash(key)}"'  # 使用key的哈希值作为ETag
+        }
+        return Response(
+            buffer.tobytes(), 
+            media_type="image/jpeg",
+            headers=headers
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/api/screenshot")
+def screenshot():
+    from ..context import device
+    img = device.screenshot()
+    buff = cv2.imencode('.png', img)[1].tobytes()
+    return Response(buff, media_type="image/png")
+
+class RunCodeRequest(BaseModel):
+    code: str
+
+@app.post("/api/code/run")
+async def run_code(request: RunCodeRequest):
+    event = asyncio.Event()
+    stdout = StringIO()
+    code = f"from kotonebot import *\n" + request.code
+    result = {}
+    def _runner():
+        nonlocal result
+        from kotonebot.backend.context import vars as context_vars
+        try:
+            with manual_context():
+                global_vars = dict(vars(kotonebot.backend.context))
+                with redirect_stdout(stdout):
+                    exec(code, global_vars)
+            result = {"status": "ok", "result": stdout.getvalue()}
+        except (Exception) as e:
+            result = {"status": "error", "result": stdout.getvalue(), "message": str(e), "traceback": traceback.format_exc()}
+        except KeyboardInterrupt as e:
+            result = {"status": "error", "result": stdout.getvalue(), "message": str(e), "traceback": traceback.format_exc()}
+        finally:
+            context_vars.flow.clear_interrupt()
+        event.set()
+    threading.Thread(target=_runner, daemon=True).start()
+    await event.wait()
+    return result
+
+@app.get("/api/code/stop")
+async def stop_code():
+    from kotonebot.backend.context import vars
+    vars.flow.request_interrupt()
+    while vars.flow.is_interrupted:
+        await asyncio.sleep(0.1)
+    return {"status": "ok"}
+
+@app.get("/api/fs/list_dir")
+def list_dir(path: str) -> list[File]:
+    result = []
+    for item in Path(path).iterdir():
+        result.append(File(
+            name=item.name,
+            full_path=str(item),
+            type="file" if item.is_file() else "dir"
+        ))
+    return result
+
+@app.get("/api/resources/autocomplete")
+def autocomplete(class_path: str) -> list[str]:
+    from kotonebot.kaa.tasks import R
+    class_names = class_path.split(".")[:-1]
+    target_class = R
+    # 定位到目标类
+    for name in class_names:
+        target_class = getattr(target_class, name, None)
+        if target_class is None:
+            return []
+    # 获取目标类的所有属性
+    attrs = [attr for attr in dir(target_class) if not attr.startswith("_")]
+    filtered_attrs = []
+    for attr in attrs:
+        if inspect.isclass(getattr(target_class, attr)):
+            filtered_attrs.append(attr)
+        elif isinstance(getattr(target_class, attr), (Image, HintBox)):
+            filtered_attrs.append(attr)  
+    attrs = filtered_attrs
+    # 排序
+    attrs.sort(key=lambda x: fuzz.ratio(x, class_path), reverse=True)
+    return attrs
+
+@app.get("/api/ping")
+async def ping():
+    return {"status": "ok"}
+
+message_queue = deque()
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    try:
+        while True:
+            if len(message_queue) > 0:
+                message = message_queue.popleft()
+                await websocket.send_json(message)
+            await asyncio.sleep(0.1)
+    except:
+        await websocket.close()
+
+def send_ws_message(title: str, image: list[str], text: str = '', callstack: list[WSCallstack] = [], wait: bool = False):
+    """发送 WebSocket 消息"""
+    message = WSMessage(
+        type="visual",
+        data=WSMessageData(
+            image=WSImage(type="memory", value=image),
+            name=title,
+            details=text,
+            timestamp=int(time.time() * 1000),
+            callstack=callstack
+        )
+    )
+    message_queue.append(message.dict())
+    if wait:
+        while len(message_queue) > 0:
+            time.sleep(0.3)
+
+
+thread = None
+def start_server():
+    global thread
+    def run_server():
+        uvicorn.run(app, host="127.0.0.1", port=8000, log_level='critical' if debug_vars.debug.hide_server_log else None)
+    if thread is None:
+        thread = threading.Thread(target=run_server, daemon=True)
+        thread.start()
+
+def wait_message_all_done():
+    global thread
+    def _wait():
+        while len(message_queue) > 0:
+            time.sleep(0.1)
+    if thread is not None:
+        threading.Thread(target=_wait, daemon=True).start()
+
+if __name__ == "__main__":
+    debug_vars.debug.hide_server_log = False
+    process = subprocess.Popen(["pylsp", "--port", "5479", "--ws"])
+    print("LSP started. PID=", process.pid)
+    uvicorn.run(app, host="127.0.0.1", port=8000, log_level='critical' if debug_vars.debug.hide_server_log else None)
+    process.kill()