Files

94 lines
3.1 KiB
Python

"""Read-only screen adapter for plugin-controlled computers. Never fabricates images."""
import base64
from io import BytesIO
from PIL import Image
from artemis.drivers.base import BaseDeviceDriver, ScreenData
class PluginScreenDriver(BaseDeviceDriver):
def __init__(self, ctx):
self.ctx = ctx
@property
def device_id(self):
return self.ctx.device.device_id
@property
def screen_size(self):
return self.ctx.device.device_width, self.ctx.device.device_height
async def connect(self):
return None
async def disconnect(self):
await self.ctx.plugins.aclose()
async def get_screen_data(self, skip_settling=False):
runtime = self.ctx.plugins
candidates = [
(pid, package) for pid, package in runtime.packages.items() if package[0].observation
]
if not candidates:
raise ValueError("电脑任务没有可用的观察插件")
pid, (manifest, _, _) = candidates[0]
loaded = await runtime.load(pid)
if loaded.status != "success":
raise RuntimeError(loaded.text())
spec = manifest.observation
tool_name, arguments = runtime._observation_calls.get(pid, (spec.tool, spec.arguments))
name = next(
(
name
for name, (owner, tool) in runtime.bindings.items()
if owner == pid and tool.name == tool_name
),
None,
)
if not name:
raise ValueError("电脑观察工具未授权或未注册")
outcome = await runtime.call(name, arguments)
if outcome.status != "success":
raise RuntimeError(outcome.text())
image = next((b for b in outcome.content if b.get("type") == "image"), None)
raw = base64.b64decode(image["data"], validate=True) if image else b""
width, height = 0, 0
if raw:
with Image.open(BytesIO(raw)) as picture:
width, height = picture.size
picture.verify()
return ScreenData(
screenshot_bytes=raw,
screenshot_base64=base64.b64encode(raw).decode(),
width=width,
height=height,
platform="computer",
observation_text=f"来源:{manifest.name} {manifest.version} / {tool_name}\n{outcome.text()}",
)
async def _unsupported(self, *args, **kwargs):
raise RuntimeError("此任务通过已授权的电脑插件操作,请使用插件工具")
tap = _unsupported
long_press = _unsupported
swipe = _unsupported
swipe_direction = _unsupported
input_text = _unsupported
press_key = _unsupported
launch_app = _unsupported
stop_app = _unsupported
get_current_package = _unsupported
execute_shell = _unsupported
start_video_recording = _unsupported
stop_video_recording = _unsupported
class PluginReadSession:
def __init__(self, ctx):
self.driver = PluginScreenDriver(ctx)
async def screenshot_b64(self):
return (await self.driver.get_screen_data()).screenshot_base64