Update tja2.py

libs/tja2.py

@@ -1,6 +1,7 @@
 import bisect
 import hashlib
 from dataclasses import dataclass, field, fields
+from typing_extensions import Optional
 
 from libs.tja import NoteList, ScrollType, TJAParser, TimelineObject, get_ms_per_measure
 
@@ -116,89 +117,109 @@ class Balloon(Note):
         hash_string = str(field_values)
         return hash_string.encode('utf-8')
 
+@dataclass
+class ParserState:
+    time_signature: float = 4/4
+    bpm: float = 120
+    bpmchange_last_bpm: float = 120
+    scroll_x_modifier: float = 1
+    scroll_y_modifier: float = 0
+    scroll_type: ScrollType = ScrollType.NMSCROLL
+    barline_display: bool = True
+    curr_note_list: list[Note | Drumroll | Balloon] = []
+    curr_draw_list: list[Note | Drumroll | Balloon] = []
+    curr_bar_list: list[Note] = []
+    curr_timeline: list[TimelineObject] = []
+    index: int = 0
+    balloons: list[int] = []
+    balloon_index: int = 0
+    prev_note: Optional[Note] = None
+    barline_added: bool = False
+
+
 class TJAParser2(TJAParser):
     def _build_command_registry(self):
-            """Auto-discover command handlers based on naming convention."""
-            registry = {}
-            for name in dir(self):
-                if name.startswith('handle_'):
-                    cmd_name = '#' + name[7:].upper()
-                    registry[cmd_name] = getattr(self, name)
-            return registry
+        """Auto-discover command handlers based on naming convention."""
+        registry = {}
+        for name in dir(self):
+            if name.startswith('handle_'):
+                cmd_name = '#' + name[7:].upper()
+                registry[cmd_name] = getattr(self, name)
+        return registry
 
-    def handle_measure(self, part: str, state: dict):
+    def handle_measure(self, part: str, state: ParserState):
         divisor = part.find('/')
-        state["time_signature"] = float(part[9:divisor]) / float(part[divisor+1:])
+        state.time_signature = float(part[9:divisor]) / float(part[divisor+1:])
 
-    def handle_scroll(self, part: str, state: dict):
+    def handle_scroll(self, part: str, state: ParserState):
         scroll_value = part[7:]
         if 'i' in scroll_value:
             normalized = scroll_value.replace('.i', 'j').replace('i', 'j')
             normalized = normalized.replace(',', '')
             c = complex(normalized)
-            state["scroll_x_modifier"] = c.real
-            state["scroll_y_modifier"] = c.imag
+            state.scroll_x_modifier = c.real
+            state.scroll_y_modifier = c.imag
         else:
-            state["scroll_x_modifier"] = float(scroll_value)
-            state["scroll_y_modifier"] = 0.0
+            state.scroll_x_modifier = float(scroll_value)
+            state.scroll_y_modifier = 0.0
 
-    def handle_bpmchange(self, part: str, state: dict):
+    def handle_bpmchange(self, part: str, state: ParserState):
         parsed_bpm = float(part[11:])
-        if state["scroll_type"] == ScrollType.BMSCROLL or state["scroll_type"] == ScrollType.HBSCROLL:
+        if state.scroll_type == ScrollType.BMSCROLL or state.scroll_type == ScrollType.HBSCROLL:
             # Do not modify bpm, it needs to be changed live by bpmchange
-            bpmchange = parsed_bpm / state["bpmchange_last_bpm"]
-            state["bpmchange_last_bpm"] = parsed_bpm
+            bpmchange = parsed_bpm / state.bpmchange_last_bpm
+            state.bpmchange_last_bpm = parsed_bpm
 
             bpmchange_timeline = TimelineObject()
             bpmchange_timeline.hit_ms = self.current_ms
             bpmchange_timeline.bpmchange = bpmchange
-            state["curr_timeline"].append(bpmchange_timeline)
+            state.curr_timeline.append(bpmchange_timeline)
         else:
             timeline_obj = TimelineObject()
             timeline_obj.hit_ms = self.current_ms
             timeline_obj.bpm = parsed_bpm
-            state["bpm"] = parsed_bpm
-            state["curr_timeline"].append(timeline_obj)
+            state.bpm = parsed_bpm
+            state.curr_timeline.append(timeline_obj)
 
-    def add_bar(self, state: dict):
+    def add_bar(self, state: ParserState):
         bar_line = Note()
 
         bar_line.hit_ms = self.current_ms
         bar_line.type = 0
-        bar_line.display = state["barline_display"]
-        bar_line.bpm = state["bpm"]
-        bar_line.scroll_x = state["scroll_x_modifier"]
-        bar_line.scroll_y = state["scroll_y_modifier"]
+        bar_line.display = state.barline_display
+        bar_line.bpm = state.bpm
+        bar_line.scroll_x = state.scroll_x_modifier
+        bar_line.scroll_y = state.scroll_y_modifier
 
-        if state["barline_added"]:
+        if state.barline_added:
             bar_line.display = False
 
         return bar_line
 
-    def add_note(self, item: str, state: dict):
+    def add_note(self, item: str, state: ParserState):
         note = Note()
         note.hit_ms = self.current_ms
         note.display = True
         note.type = int(item)
-        note.index = state["index"]
-        note.bpm = state["bpm"]
-        note.scroll_x = state["scroll_x_modifier"]
-        note.scroll_y = state["scroll_y_modifier"]
+        note.index = state.index
+        note.bpm = state.bpm
+        note.scroll_x = state.scroll_x_modifier
+        note.scroll_y = state.scroll_y_modifier
 
         if item in {'5', '6'}:
             note = Drumroll(note)
             note.color = 255
         elif item in {'7', '9'}:
-            state["balloon_index"] += 1
-            if state["balloons"] is None:
+            state.balloon_index += 1
+            if state.balloons is None:
                 raise Exception("Balloon note found, but no count was specified")
             if item == '9':
                 note = Balloon(note, is_kusudama=True)
             else:
                 note = Balloon(note)
-            note.count = 1 if not state["balloons"] else state["balloons"].pop(0)
+            note.count = 1 if not state.balloons else state.balloons.pop(0)
         elif item == '8':
-            if state["prev_note"] is None:
+            if state.prev_note is None:
                 raise ValueError("No previous note found")
 
         return note
@@ -209,32 +230,23 @@ class TJAParser2(TJAParser):
         master_notes = NoteList()
         notes = self.data_to_notes(diff)
 
-        state = {
-            'time_signature': 4/4,
-            'bpm': self.metadata.bpm,
-            'scroll_x_modifier': 1,
-            'scroll_y_modifier': 0,
-            'scroll_type': ScrollType.NMSCROLL,
-            'bpmchange_last_bpm': self.metadata.bpm,
-            'barline_display': True,
-            'curr_note_list': master_notes.play_notes,
-            'curr_draw_list': master_notes.draw_notes,
-            'curr_bar_list': master_notes.bars,
-            'curr_timeline': master_notes.timeline,
-            'index': 0,
-            'balloons': self.metadata.course_data[diff].balloon.copy(),
-            'balloon_index': 0,
-            'prev_note': None,
-            'barline_added': False
-        }
+        state = ParserState()
+        state.bpm = self.metadata.bpm
+        state.bpmchange_last_bpm = self.metadata.bpm
+        state.balloons = self.metadata.course_data[diff].balloon.copy()
+        state.curr_note_list = master_notes.play_notes
+        state.curr_draw_list = master_notes.draw_notes
+        state.curr_bar_list = master_notes.bars
+        state.curr_timeline = master_notes.timeline
+
         init_bpm = TimelineObject()
         init_bpm.hit_ms = self.current_ms
-        init_bpm.bpm = state['bpm']
-        state['curr_timeline'].append(init_bpm)
+        init_bpm.bpm = state.bpm
+        state.curr_timeline.append(init_bpm)
 
         for bar in notes:
             bar_length = sum(len(part) for part in bar if '#' not in part)
-            state['barline_added'] = False
+            state.barline_added = False
 
             for part in bar:
                 if part.startswith('#'):
@@ -246,11 +258,11 @@ class TJAParser2(TJAParser):
                 elif len(part) > 0 and not part[0].isdigit():
                     continue
 
-                ms_per_measure = get_ms_per_measure(state["bpm"], state["time_signature"])
+                ms_per_measure = get_ms_per_measure(state.bpm, state.time_signature)
 
                 bar = self.add_bar(state)
-                state["curr_bar_list"].append(bar)
-                state["barline_added"] = True
+                state.curr_bar_list.append(bar)
+                state.barline_added = True
 
                 if len(part) == 0:
                     self.current_ms += ms_per_measure
@@ -266,10 +278,10 @@ class TJAParser2(TJAParser):
                     note = self.add_note(item, state)
 
                     self.current_ms += increment
-                    state["curr_note_list"].append(note)
-                    state["curr_draw_list"].append(note)
-                    self.get_moji(state["curr_note_list"], ms_per_measure)
-                    state["index"] += 1
-                    state["prev_note"] = note
+                    state.curr_note_list.append(note)
+                    state.curr_draw_list.append(note)
+                    self.get_moji(state.curr_note_list, ms_per_measure)
+                    state.index += 1
+                    state.prev_note = note
 
         return master_notes, [master_notes], [master_notes], [master_notes]