fix typing \o/

src/writer/app_runner.py

@@ -1,7 +1,7 @@
 import asyncio
 import concurrent.futures
-import importlib.util
 import glob
+import importlib.util
 import logging
 import logging.handlers
 import multiprocessing
@@ -44,6 +44,7 @@
     InitSessionRequestPayload,
     InitSessionResponsePayload,
     ServeMode,
+    SourceFilesDirectory,
     StateContentRequest,
     StateContentResponsePayload,
     StateEnquiryRequest,
@@ -105,7 +106,7 @@ def __init__(self,
         self.app_path = app_path
         self.mode = mode
         self.run_code = run_code
-        self.source_files: Dict = {}
+        self.source_files: SourceFilesDirectory = {"children": {}, "type": "directory"}
         self.bmc_components = bmc_components
         self.is_app_process_server_ready = is_app_process_server_ready
         self.is_app_process_server_failed = is_app_process_server_failed 
@@ -607,7 +608,7 @@ def __init__(self, app_path: str, mode: str):
         self.client_conn: Optional[multiprocessing.connection.Connection] = None
         self.app_process: Optional[AppProcess] = None
         self.run_code: Optional[str] = None
-        self.source_files: Dict = {}
+        self.source_files: SourceFilesDirectory = {"children": {}, "type": "directory"}
         self.bmc_components: Optional[Dict] = None
         self.is_app_process_server_ready = multiprocessing.Event()
         self.is_app_process_server_failed = multiprocessing.Event()
@@ -671,7 +672,7 @@ def load(self) -> None:
         # parent pid and pid.
         self._subscribe_terminal_signal()
 
-    def _build_source_files(self) -> Dict:
+    def _build_source_files(self) -> SourceFilesDirectory:
         """
         Build a file tree as `Dict` wherein the key represent the filename. The value is a `Dict` with a `type` as:
 
@@ -683,27 +684,31 @@ def _build_source_files(self) -> Dict:
         >>> {'type': 'directory', 'children': {'README.md': {'type': 'file', 'content': 'This app w', 'complete': False}}}
         """
         files = glob.glob(os.path.join(self.app_path, '**', "*"), recursive=True)
-        file_tree = {
+        file_tree: SourceFilesDirectory = {
             "type": "directory",
             "children": {}
         }
 
         for file in files:
             relative_path = os.path.relpath(file, self.app_path)
             parts = relative_path.split(os.sep)
-            current_level = file_tree
+            current_level: SourceFilesDirectory = file_tree
 
             for part in parts:
                 if os.path.isdir(os.path.join(self.app_path, *parts[:parts.index(part) + 1])):
-                    if part not in current_level['children']:
-                        current_level['children'][part] = { "type": "directory", "children": {} }
-                    current_level = current_level['children'][part]
+                    if part not in current_level["children"]:
+                        current_level["children"][part] = { "type": "directory", "children": {} }
+
+                    next_level = current_level["children"][part]
+
+                    if next_level["type"] == "directory":
+                        current_level = next_level
                 else:
                     try:
                         content = self.load_persisted_script(relative_path)
                         # limit only the first 100 characters to limit bandwidth usage, the rest will be lazy loaded
                         exerpt = content[0:100]
-                        current_level['children'][part] = {
+                        current_level["children"][part] = {
                             "type": "file",
                             "content": exerpt,
                             "complete": exerpt == content,

src/writer/ss_types.py

@@ -33,6 +33,19 @@ class AbstractTemplate(BaseModel):
     baseType: str
     writer: Dict
 
+
+class SourceFilesFile(TypedDict):
+    type: Literal["file"]
+    complete: Optional[bool]
+    content: str
+
+
+class SourceFilesDirectory(TypedDict):
+    type: Literal["directory"]
+    children: Dict[str, 'SourceFiles']
+
+SourceFiles = Union[SourceFilesFile, SourceFilesDirectory]
+
 # Web server models
 
 
@@ -59,7 +72,7 @@ class InitResponseBodyRun(InitResponseBody):
 class InitResponseBodyEdit(InitResponseBody):
     mode: Literal["edit"]
     runCode: Optional[str] = None
-    sourceFiles: Dict = {}
+    sourceFiles: SourceFilesDirectory = {"type": "directory", "children": {}}
 
 
 class WriterWebsocketIncoming(BaseModel):
@@ -232,3 +245,4 @@ class WorkflowExecutionLog(BaseModel):
 
 class WriterConfigurationError(ValueError):
     pass
+