ksp_client.py 2.11 KB
Newer Older
Luca Cristaldi's avatar
Luca Cristaldi committed
1 2
import krpc
import zmq
Luca Cristaldi's avatar
Luca Cristaldi committed
3
import argparse
4
from threading import Event
Luca Cristaldi's avatar
Luca Cristaldi committed
5 6


7 8
def main():
    parser = argparse.ArgumentParser(description='')
9
    parser.add_argument('addr', type=str, default="*:5556", nargs="?")
10
    args = parser.parse_args()
11 12

    host, port = args.addr.split(":")
Luca Cristaldi's avatar
Luca Cristaldi committed
13
    topic = "events"
14 15 16 17 18
    states = {-1: "abort",
              0: "pre_launch",
              1: "liftoff",
              2: "stage_1_separation",
              3: "stage_2_separation"}
19
    stage_map = dict(zip(range(10), range(10, 0, -1)))
Luca Cristaldi's avatar
Luca Cristaldi committed
20

21 22 23
    context = zmq.Context()
    socket = context.socket(zmq.PUB)
    socket.bind(f"tcp://{host}:{port}")
Luca Cristaldi's avatar
Luca Cristaldi committed
24

25
    conn = krpc.connect()
26

27 28
    running = Event()
    stopping = Event()
29

30 31 32 33 34 35 36 37
    def unlock(scene):
        print(f"In GameScene {scene.name}")
        if (scene == conn.krpc.GameScene.flight):
            running.set()
            stopping.clear()
        else:
            running.clear()
            stopping.set()
Luca Cristaldi's avatar
Luca Cristaldi committed
38

39 40
    def stage_changed(new_stage):
        try:
41 42 43 44
            stage_mapped = stage_map.get(new_stage, new_stage)
            stage_name = states.get(stage_mapped, "Unknown")
            socket.send_string(f"{topic} {stage_name}")
            print(f"Sending event {stage_name} ({stage_mapped})")
45
        except Exception as e:
46 47
            print("Exception occurred in callback")
            print(e)
Luca Cristaldi's avatar
Luca Cristaldi committed
48

49 50 51
    scene = conn.add_stream(getattr, conn.krpc, 'current_game_scene')
    scene.add_callback(unlock)
    scene.start()
52

53 54 55 56 57
    while True:
        running.wait()
        running.clear()
        vessel = conn.space_center.active_vessel
        control = vessel.control
58 59 60
        vehicle_stages = max([p.stage for p in vessel.parts.all]) + 1
        stage_map = dict(zip(range(vehicle_stages + 1),
                             range(vehicle_stages, -1, -1)))
61
        stage = conn.add_stream(getattr, control, 'current_stage', scene)
62

63 64
        stage.add_callback(stage_changed)
        stage.start()
65

66 67
        stopping.wait()
        stopping.clear()
68
        socket.send_string("events abort")
69 70 71
        print("Removing callback...")
        stage.remove_callback(stage_changed)
        stage.remove()
72

Luca Cristaldi's avatar
Luca Cristaldi committed
73

74 75
if __name__ == "__main__":
    main()