Passed
Push — master ( 6e56fb...e44069 )
by manny
01:57
created

racetime_obs.RacetimeObs.process_ws_message()   A

Complexity

Conditions 5

Size

Total Lines 14
Code Lines 14

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
eloc 14
dl 0
loc 14
rs 9.2333
c 0
b 0
f 0
cc 5
nop 3
1
import json
2
import websockets
3
from websockets.client import WebSocketClientProtocol
4
import obspython as obs
5
from datetime import datetime, timedelta, timezone
6
import racetime_client
7
from models.race import Entrant, Race, race_from_dict
8
import asyncio
9
from threading import Thread
10
import websockets
11
import dateutil
12
import logging
13
from helpers.LogFormatter import LogFormatter
14
from helpers.obs_context_manager import source_ar, source_list_ar, data_ar
15
from gadgets.timer import Timer
16
from gadgets.coop import Coop
17
from gadgets.qualifier import Qualifier
18
19
class RacetimeObs():
20
    logger = logging.Logger("racetime-obs")
21
    race: Race = None
22
    selected_race = ""
23
    check_race_updates = False
24
    race_changed = False
25
    full_name = ""
26
    category = ""
27
    timer = Timer()
28
    coop = Coop()
29
    qualifier = Qualifier()
30
31
    def update_sources(self):
32
        if self.race is not None:
33
            if self.timer.enabled:
34
                color, time = self.timer.get_timer_text(self.race, self.full_name)
35
                self.set_source_text(self.timer.source_name, time, color)
36
            if self.coop.enabled:
37
                self.set_source_text(self.coop.source_name, self.coop.text, None)
38
                self.set_source_text(self.coop.label_source_name, self.coop.label_text, None)
39
            if self.qualifier.enabled:
40
                self.set_source_text(self.qualifier.qualifier_par_source, self.qualifier.qualifier_par_text, None)
41
                self.set_source_text(self.qualifier.qualifier_score_source, self.qualifier.entrant_score, None)
42
            pass
43
44
    def race_update_thread(self):
45
        self.logger.debug("starting race update")
46
        race_event_loop = asyncio.new_event_loop()
47
        race_event_loop.run_until_complete(self.race_updater())
48
        race_event_loop.run_forever()
49
50
51
    async def race_updater(self):
52
        headers = {
53
            'User-Agent': "oro-obs-bot_alpha"
54
        }
55
        host = "racetime.gg"
56
57
        while True:
58
            if not self.timer.enabled:
59
                await asyncio.sleep(5.0)
60
            else:
61
                if self.race is None and self.selected_race != "":
62
                    self.race = racetime_client.get_race(self.selected_race)
63
                if self.race is not None and self.race.websocket_url != "":
64
                    async with websockets.connect("wss://racetime.gg" + self.race.websocket_url, host=host, extra_headers=headers) as ws:
65
                        self.race_changed = False
66
                        self.logger.info(
67
                            f"connected to websocket: {self.race.websocket_url}")
68
                        await self.process_messages(ws)
69
            await asyncio.sleep(5.0)
70
71
72
    async def process_messages(self, ws: WebSocketClientProtocol):
73
        last_pong = datetime.now(timezone.utc)
74
        while True:
75
            try:
76
                if self.race_changed:
77
                    self.logger.info("new race selected")
78
                    self.race_changed = False
79
                    break
80
                message = await asyncio.wait_for(ws.recv(), 5.0)
81
                self.logger.info(f"received message from websocket: {message}")
82
                data = json.loads(message)
83
                last_pong = self.process_ws_message(data, last_pong)
84
            except asyncio.TimeoutError:
85
                if datetime.now(timezone.utc) - last_pong > timedelta(seconds=20):
86
                    await ws.send(json.dumps({"action": "ping"}))
87
            except websockets.ConnectionClosed:
88
                self.logger.error(f"websocket connection closed")
89
                self.race = None
90
                break
91
92
    def process_ws_message(self, data, last_pong):
93
        if data.get("type") == "race.data":
94
            r = race_from_dict(data.get("race"))
95
            self.logger.debug(f"race data parsed: {r}")
96
            self.logger.debug(f"current race is {self.race}")
97
            if r is not None and r.version > self.race.version:
98
                self.race = r
99
                self.logger.debug(f"self.race is {self.race}")
100
                self.coop.update_coop_text(self.race, self.full_name)
101
                self.qualifier.update_qualifier_text(self.race, self.full_name)
102
        elif data.get("type") == "pong":
103
            last_pong = dateutil.parser.parse(data.get("date"))
104
            pass
105
        return last_pong
106
107
108
    def update_logger(self, enabled: bool, log_to_file: bool, log_file: str, level: str):
109
        self.logger.disabled = not enabled
110
        self.logger.handlers = []
111
        handler = logging.StreamHandler()
112
        if log_to_file:
113
            try:
114
                handler = logging.FileHandler(log_file)
115
            except:
116
                self.logger.error(f"Unable to open {log_file}")
117
        elif level == "Debug":
118
            handler.setLevel(logging.DEBUG)
119
        elif level == "Info":
120
            handler.setLevel(logging.INFO)
121
        else:
122
            handler.setLevel(logging.ERROR)
123
        handler.setFormatter(LogFormatter())
124
        self.logger.addHandler(handler)
125
126
    @staticmethod
127
    def fill_source_list(p):
128
        obs.obs_property_list_clear(p)
129
        obs.obs_property_list_add_string(p, "", "")
130
        with source_list_ar() as sources:
131
            if sources is not None:
132
                for source in sources:
133
                    source_id = obs.obs_source_get_unversioned_id(source)
134
                    if source_id == "text_gdiplus" or source_id == "text_ft2_source":
135
                        name = obs.obs_source_get_name(source)
136
                        obs.obs_property_list_add_string(p, name, name)
137
138
    def fill_race_list(self, race_list, category_list):
139
        obs.obs_property_list_clear(race_list)
140
        obs.obs_property_list_clear(category_list)
141
        obs.obs_property_list_add_string(category_list, "All", "All")
142
143
        obs.obs_property_list_add_string(race_list, "", "")
144
        races = racetime_client.get_races()
145
        if races is not None:
146
            categories = []
147
            for race in races:
148
                if self.category == "" or self.category == "All" or race.category.name == self.category:
149
                    obs.obs_property_list_add_string(
150
                        race_list, race.name, race.name)
151
                if not race.category.name in categories:
152
                    categories.append(race.category.name)
153
                    obs.obs_property_list_add_string(
154
                        category_list, race.category.name, race.category.name)
155
156
157
    def fill_coop_entrant_lists(self, props):
158
        self.fill_entrant_list(obs.obs_properties_get(props, "coop_partner"))
159
        self.fill_entrant_list(obs.obs_properties_get(props, "coop_opponent1"))
160
        self.fill_entrant_list(obs.obs_properties_get(props, "coop_opponent2"))
161
162
163
    def fill_entrant_list(self, entrant_list):
164
        obs.obs_property_list_clear(entrant_list)
165
        obs.obs_property_list_add_string(entrant_list, "", "")
166
        if self.race is not None:
167
            for entrant in self.race.entrants:
168
                obs.obs_property_list_add_string(
169
                    entrant_list, entrant.user.full_name, entrant.user.full_name)
170
    
171
    
172
    # copied and modified from scripted-text.py by UpgradeQ
173
174
    @staticmethod
175
    def set_source_text(source_name: str, text: str, color: int):
176
        with source_ar(source_name) as source, data_ar() as settings:
177
            obs.obs_data_set_string(settings, "text", text)
178
            source_id = obs.obs_source_get_unversioned_id(source)
179
            if color is not None:
180
                if source_id == "text_gdiplus":
181
                    obs.obs_data_set_int(settings, "color", color)  # colored text
182
183
                else:  # freetype2,if taken from user input it should be reversed for getting correct color
184
                    number = "".join(hex(color)[2:])
185
                    color = int("0xff" f"{number}", base=16)
186
                    obs.obs_data_set_int(settings, "color1", color)
187
            
188
            obs.obs_source_update(source, settings)
189
190
191
rtgg_obs = RacetimeObs()
192
193
# ------------------------------------------------------------
194
195
196
def script_description():
197
    return "<center><p>Select a text source to use as your timer and enter your full " + \
198
        "username on racetime.gg  (including discriminator). This only needs " + \
199
        "to be done once.\n\nThen select the race room each race you join and " + \
200
        "stop worrying about whether you started your timer or not.<hr/></p>"
201
202
203
def script_load(settings):
204
    rtgg_obs.timer.use_podium_colors = obs.obs_data_get_bool(settings, "use_podium")
205
206
    race_update_t = Thread(target=rtgg_obs.race_update_thread)
207
    race_update_t.daemon = True
208
    race_update_t.start()
209
210
211
def script_save(settings):
212
    obs.obs_data_set_bool(settings, "use_podium", rtgg_obs.timer.use_podium_colors)
213
214
def script_update(settings):
215
216
    obs.timer_remove(rtgg_obs.update_sources)
217
218
    rtgg_obs.update_logger(obs.obs_data_get_bool(settings, "enable_log"),
219
        obs.obs_data_get_bool(settings, "log_to_file"),
220
        obs.obs_data_get_string(settings, "log_file"), 
221
        obs.obs_data_get_string(settings, "log_level"))
222
223
    rtgg_obs.full_name = obs.obs_data_get_string(settings, "username")
224
225
    rtgg_obs.timer.source_name = obs.obs_data_get_string(settings, "source")
226
227
    rtgg_obs.selected_race = obs.obs_data_get_string(settings, "race")
228
    rtgg_obs.category = obs.obs_data_get_string(settings, "category_filter")
229
230
    rtgg_obs.timer.use_podium_colors = obs.obs_data_get_bool(settings, "use_podium")
231
    rtgg_obs.timer.pre_color = obs.obs_data_get_int(settings, "pre_color")
232
    rtgg_obs.timer.first_color = obs.obs_data_get_int(settings, "first_color")
233
    rtgg_obs.timer.second_color = obs.obs_data_get_int(settings, "second_color")
234
    rtgg_obs.timer.third_color = obs.obs_data_get_int(settings, "third_color")
235
    rtgg_obs.timer.racing_color = obs.obs_data_get_int(settings, "racing_color")
236
    rtgg_obs.timer.finished_color = obs.obs_data_get_int(settings, "finished_color")
237
238
    if rtgg_obs.timer.source_name != "" and rtgg_obs.selected_race != "":
239
        obs.timer_add(rtgg_obs.update_sources, 100)
240
        rtgg_obs.timer.enabled = True
241
    else:
242
        rtgg_obs.timer.enabled = False
243
    rtgg_obs.logger.debug(f"timer.enabled is {rtgg_obs.timer.enabled}")
244
    rtgg_obs.logger.debug(f"timer.source_name is {rtgg_obs.timer.source_name}")
245
    rtgg_obs.logger.debug(f"selected_race is {rtgg_obs.selected_race}")
246
247
    rtgg_obs.coop.enabled = obs.obs_data_get_bool(settings, "use_coop")
248
    rtgg_obs.coop.partner = obs.obs_data_get_string(settings, "coop_partner")
249
    rtgg_obs.coop.opponent1 = obs.obs_data_get_string(settings, "coop_opponent1")
250
    rtgg_obs.coop.opponent2 = obs.obs_data_get_string(settings, "coop_opponent2")
251
    rtgg_obs.coop.source_name = obs.obs_data_get_string(settings, "coop_source")
252
    rtgg_obs.coop.label_source_name = obs.obs_data_get_string(settings, "coop_label")
253
254
    rtgg_obs.qualifier.enabled = obs.obs_data_get_bool(settings, "use_qualifier")
255
    rtgg_obs.qualifier.qualifier_cutoff = obs.obs_data_get_int(settings, "qualifier_cutoff")
256
    rtgg_obs.qualifier.qualifier_par_source = obs.obs_data_get_string(
257
        settings, "qualifier_par_source")
258
    rtgg_obs.qualifier.qualifier_score_source = obs.obs_data_get_string(
259
        settings, "qualifier_score_source")
260
261
262
def script_defaults(settings):
263
    obs.obs_data_set_default_string(settings, "race_info", "Race info")
264
    obs.obs_data_set_default_string(settings, "race", "")
265
266
def script_properties():
267
    props = obs.obs_properties_create()
268
269
    setup_group = obs.obs_properties_create()
270
    obs.obs_properties_add_group(
271
        props, "initial_setup", "Initial setup - Check to make changes", obs.OBS_GROUP_CHECKABLE, setup_group)
272
    p = obs.obs_properties_add_list(
273
        setup_group, "source", "Text Source", obs.OBS_COMBO_TYPE_EDITABLE, obs.OBS_COMBO_FORMAT_STRING)
274
    rtgg_obs.fill_source_list(p)
275
    obs.obs_properties_add_text(
276
        setup_group, "username", "Username", obs.OBS_TEXT_DEFAULT)
277
    logging = obs.obs_properties_add_bool(
278
        setup_group, "enable_log", "Enable logging")
279
    log_levels = obs.obs_properties_add_list(
280
        setup_group, "log_level", "Log lever", obs.OBS_COMBO_TYPE_LIST, obs.OBS_COMBO_FORMAT_STRING)
281
    obs.obs_property_list_add_string(log_levels, "Error", "Error")
282
    obs.obs_property_list_add_string(log_levels, "Debug", "Debug")
283
    obs.obs_property_list_add_string(log_levels, "Info", "Info")
284
    obs.obs_property_set_long_description(
285
        logging, "Generally, only log errors unless you are developing or are trying to find a specific problem.")
286
    obs.obs_properties_add_bool(setup_group, "log_to_file", "Log to file?")
287
    #obs.obs_property_set_modified_callback(p, log_to_file_toggled)
288
    obs.obs_properties_add_path(
289
        setup_group, "log_file", "Log File", obs.OBS_PATH_FILE_SAVE, "Text files(*.txt)", None)
290
291
    category_list = obs.obs_properties_add_list(
292
        props, "category_filter", "Filter by Category", obs.OBS_COMBO_TYPE_LIST, obs.OBS_COMBO_FORMAT_STRING)
293
    race_list = obs.obs_properties_add_list(
294
        props, "race", "Race", obs.OBS_COMBO_TYPE_LIST, obs.OBS_COMBO_FORMAT_STRING)
295
    obs.obs_property_set_modified_callback(race_list, new_race_selected)
296
    obs.obs_property_set_modified_callback(
297
        category_list, new_category_selected)
298
299
    p = obs.obs_properties_add_text(
300
        props, "race_info", "Race Desc", obs.OBS_TEXT_MULTILINE)
301
    obs.obs_property_set_enabled(p, False)
302
303
    refresh = obs.obs_properties_add_button(
304
        props, "button", "Refresh", lambda *props: None)
305
    obs.obs_property_set_modified_callback(refresh, refresh_pressed)
306
307
    p = obs.obs_properties_add_bool(
308
        props, "use_podium", "Use custom color for podium finishes?")
309
    obs.obs_property_set_modified_callback(p, podium_toggled)
310
311
    podium_group = obs.obs_properties_create()
312
    obs.obs_properties_add_group(
313
        props, "podium_group", "Podium Colors", obs.OBS_GROUP_NORMAL, podium_group)
314
    obs.obs_property_set_visible(obs.obs_properties_get(
315
        props, "podium_group"), rtgg_obs.timer.use_podium_colors)
316
317
    obs.obs_properties_add_color(podium_group, "pre_color", "Pre-race:")
318
    obs.obs_properties_add_color(podium_group, "racing_color", "Still racing:")
319
    obs.obs_properties_add_color(podium_group, "first_color", "1st place:")
320
    obs.obs_properties_add_color(podium_group, "second_color", "2nd place:")
321
    obs.obs_properties_add_color(podium_group, "third_color", "3rd place:")
322
    obs.obs_properties_add_color(
323
        podium_group, "finished_color", "After podium:")
324
325
    p = obs.obs_properties_add_bool(
326
        props, "use_coop", "Display coop information?")
327
    obs.obs_property_set_modified_callback(p, coop_toggled)
328
329
    coop_group = obs.obs_properties_create()
330
    obs.obs_properties_add_group(
331
        props, "coop_group", "Co-op Mode", obs.OBS_GROUP_NORMAL, coop_group)
332
    obs.obs_property_set_visible(
333
        obs.obs_properties_get(props, "coop_group"), rtgg_obs.coop.enabled)
334
    p = obs.obs_properties_add_list(
335
        coop_group, "coop_partner", "Co-op Partner", obs.OBS_COMBO_TYPE_LIST, obs.OBS_COMBO_FORMAT_STRING)
336
    p = obs.obs_properties_add_list(
337
        coop_group, "coop_opponent1", "Co-op Opponent 1", obs.OBS_COMBO_TYPE_LIST, obs.OBS_COMBO_FORMAT_STRING)
338
    p = obs.obs_properties_add_list(
339
        coop_group, "coop_opponent2", "Co-op Opponent 2", obs.OBS_COMBO_TYPE_LIST, obs.OBS_COMBO_FORMAT_STRING)
340
    rtgg_obs.fill_coop_entrant_lists(props)
341
    p = obs.obs_properties_add_list(coop_group, "coop_source", "Coop Text Source",
342
                                    obs.OBS_COMBO_TYPE_EDITABLE, obs.OBS_COMBO_FORMAT_STRING)
343
    obs.obs_property_set_long_description(
344
        p, "This text source will display the time that the last racer needs to finish for their team to win")
345
    rtgg_obs.fill_source_list(p)
346
    p = obs.obs_properties_add_list(coop_group, "coop_label", "Coop Label Text Source",
347
                                    obs.OBS_COMBO_TYPE_EDITABLE, obs.OBS_COMBO_FORMAT_STRING)
348
    obs.obs_property_set_long_description(
349
        p, "This text source will be use to display a label such as \'<PartnerName> needs to finish before\' based on who the last racer is")
350
    rtgg_obs.fill_source_list(p)
351
352
    p = obs.obs_properties_add_bool(
353
        props, "use_qualifier", "Display race results as tournament qualifier?")
354
    obs.obs_property_set_modified_callback(p, qualifier_toggled)
355
356
    qualifier_group = obs.obs_properties_create()
357
    obs.obs_properties_add_group(
358
        props, "qualifier_group", "Qualifier Mode", obs.OBS_GROUP_NORMAL, qualifier_group)
359
    obs.obs_property_set_visible(
360
        obs.obs_properties_get(props, "qualifier_group"), rtgg_obs.qualifier.enabled)
361
    p = obs.obs_properties_add_int_slider(
362
        qualifier_group, "qualifier_cutoff", "Use Top X as par time, where X=", 3, 10, 1)
363
    p = obs.obs_properties_add_list(qualifier_group, "qualifier_par_source",
364
                                    "Qualifier Par Time Source", obs.OBS_COMBO_TYPE_EDITABLE, obs.OBS_COMBO_FORMAT_STRING)
365
    rtgg_obs.fill_source_list(p)
366
    p = obs.obs_properties_add_list(qualifier_group, "qualifier_score_source",
367
                                    "Qualifier Score Source", obs.OBS_COMBO_TYPE_EDITABLE, obs.OBS_COMBO_FORMAT_STRING)
368
    rtgg_obs.fill_source_list(p)
369
370
    return props
371
372
def refresh_pressed(props, prop, *args, **kwargs):
373
    rtgg_obs.fill_source_list(obs.obs_properties_get(props, "source"))
374
    rtgg_obs.fill_source_list(obs.obs_properties_get(props, "coop_label"))
375
    rtgg_obs.fill_source_list(obs.obs_properties_get(props, "coop_text"))
376
    rtgg_obs.fill_source_list(obs.obs_properties_get(props, "qualifier_par_source"))
377
    rtgg_obs.fill_source_list(obs.obs_properties_get(props, "qualifier_score_source"))
378
    rtgg_obs.fill_race_list(obs.obs_properties_get(props, "race"),
379
                obs.obs_properties_get(props, "category_filter"))
380
    if rtgg_obs.race is not None:
381
        rtgg_obs.coop.update_coop_text(rtgg_obs.race, rtgg_obs.full_name)
382
        rtgg_obs.qualifier.update_qualifier_text(rtgg_obs.race, rtgg_obs.full_name)
383
    return True
384
385
386
def new_race_selected(props, prop, settings):
387
    rtgg_obs.selected_race = obs.obs_data_get_string(settings, "race")
388
    r = racetime_client.get_race(rtgg_obs.selected_race)
389
    if r is not None:
390
        rtgg_obs.race = r
391
        rtgg_obs.coop.update_coop_text(rtgg_obs.race, rtgg_obs.full_name)
392
        rtgg_obs.qualifier.update_qualifier_text(rtgg_obs.race, rtgg_obs.full_name)
393
        rtgg_obs.logger.info(f"new race selected: {rtgg_obs.race}")
394
        obs.obs_data_set_default_string(settings, "race_info", r.info)
395
        rtgg_obs.fill_coop_entrant_lists(props)
396
    else:
397
        obs.obs_data_set_default_string(
398
            settings, "race_info", "Race not found")
399
400
    rtgg_obs.race_changed = True
401
    return True
402
403
404
def new_category_selected(props, prop, settings):
405
    rtgg_obs.category = obs.obs_data_get_string(settings, "category_filter")
406
    rtgg_obs.logger.info(f"new category selected: {rtgg_obs.category}")
407
    rtgg_obs.fill_race_list(obs.obs_properties_get(props, "race"), prop)
408
    return True
409
410
411
def podium_toggled(props, prop, settings):
412
    vis = obs.obs_data_get_bool(settings, "use_podium")
413
    obs.obs_property_set_visible(
414
        obs.obs_properties_get(props, "podium_group"), vis)
415
    return True
416
417
418
def coop_toggled(props, prop, settings):
419
    vis = obs.obs_data_get_bool(settings, "use_coop")
420
    obs.obs_property_set_visible(
421
        obs.obs_properties_get(props, "coop_group"), vis)
422
    return True
423
424
425
def qualifier_toggled(props, prop, settings):
426
    vis = obs.obs_data_get_bool(settings, "use_qualifier")
427
    obs.obs_property_set_visible(
428
        obs.obs_properties_get(props, "qualifier_group"), vis)
429
    return True