41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
from enum import Enum
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from simulator import EmulatorSimulator, MouseClick
|
|
|
|
|
|
class Action(Enum):
|
|
CLICK = "click"
|
|
MOUSE_DOWN = "mousedown"
|
|
MOUSE_UP = "mouseup"
|
|
MOUSE_MOVE = "mousemove"
|
|
KEY = "key"
|
|
TYPE = "type"
|
|
|
|
|
|
class Controller:
|
|
def __init__(self, vm_name: str, username: str, password: str, host: str) -> None:
|
|
self.simulator = EmulatorSimulator(vm_name=vm_name, username=username,
|
|
password=password, host=host)
|
|
|
|
def get_state(self) -> np.ndarray:
|
|
image_path = self.simulator.get_screenshot()
|
|
with Image.open(image_path) as img:
|
|
return np.array(img)
|
|
|
|
def step(self, action: Action, **action_args) -> None:
|
|
if action == Action.CLICK:
|
|
print(action_args)
|
|
self.simulator.click(**action_args)
|
|
elif action == Action.MOUSE_DOWN:
|
|
self.simulator.mousedown(**action_args)
|
|
elif action == Action.MOUSE_UP:
|
|
self.simulator.mouseup(**action_args)
|
|
elif action == Action.MOUSE_MOVE:
|
|
self.simulator.mouse_move(**action_args)
|
|
elif action == Action.KEY:
|
|
self.simulator.key(**action_args)
|
|
elif action == Action.TYPE:
|
|
self.simulator.type(**action_args) |