Fix typing

Author: solidpixel

lglpy/comms/service_gpu_profile.py

@@ -29,7 +29,7 @@
 import csv
 import json
 import os
-from typing import TypedDict, Union
+from typing import Any, Optional, TypedDict, Union
 
 from lglpy.comms.server import Message
 
@@ -73,9 +73,9 @@ def __init__(self, dir_path: str, verbose: bool = False):
         '''
         self.base_dir = dir_path
 
-        self.frame_id = None
-        self.frame_header = None
-        self.frame_data = None
+        self.frame_id: Optional[int] = None
+        self.frame_header: Optional[list[str]] = None
+        self.frame_data: Optional[list[list[str]]] = None
 
         os.makedirs(dir_path, exist_ok=True)
 
@@ -109,6 +109,10 @@ def handle_end_frame(self, message: EndFrameMessage):
         # Message contains nothing we need
         del message
 
+        assert self.frame_id is not None
+        assert self.frame_header is not None
+        assert self.frame_data is not None
+
         # Emit the CSV file
         print(f'Generating CSV for frame {self.frame_id}')
         path = os.path.join(self.base_dir, f'frame_{self.frame_id}.csv')
@@ -147,13 +151,18 @@ def create_workload_data(self, message: WorkloadMessage):
         Args:
             message: The decoded JSON.
         '''
-        columns = []
+        assert self.frame_id is not None
+        assert self.frame_header is not None
+        assert self.frame_data is not None
+
+        columns: list[str] = []
 
-        columns.append(len(self.frame_data))
+        columns.append(str(len(self.frame_data)))
         columns.append(message['type'])
+
         for counter in message['counters']:
             value = list(counter.values())[0]
-            columns.append(f'{value:0.2f}')
+            columns.append( f'{value:0.2f}')
         columns.append('|'.join(message['labels']))
 
         self.frame_data.append(columns)