Fix for all affected modules

Author: zanjonke

partial_rendering.py

@@ -116,6 +116,34 @@ def get_plain_module_render_state(plain_module: PlainModule) -> PlainModuleRende
     return pr
 
 
+def get_all_affected_modules_from_change(
+    plain_module: PlainModule,
+    plain_module_render_state: PlainModuleRenderState,
+) -> list[PlainModule]:
+    all_affected_modules = list[PlainModule]()
+
+    if plain_module_render_state.change_type == "spec_change":
+        start_module = plain_module_render_state.change
+    elif plain_module_render_state.change_type == "code_change":
+        if plain_module_render_state.change.is_module_fully_rendered():
+            start_module = plain_module.get_next_module(plain_module_render_state.change.module_name)
+        else:
+            start_module = plain_module_render_state.change
+    else:
+        raise ValueError(f"Unknown change type: {plain_module_render_state.change_type}")
+
+    affected_module = False
+    all_modules = plain_module.all_required_modules + [plain_module]
+    for module in all_modules:
+        if module.module_name == start_module.module_name:
+            affected_module = True
+
+        if affected_module:
+            all_affected_modules.append(module)
+
+    return all_affected_modules
+
+
 def get_render_choices(
     plain_module: PlainModule,
     plain_module_render_state: PlainModuleRenderState,
@@ -165,17 +193,11 @@ def get_render_choices(
         choice_idx += 1
 
     if plain_module_render_state.change:
-        all_affected_modules = list[str]()
-        affected_module = False
-        for module in plain_module.all_required_modules:
-            if module.module_name == plain_module_render_state.change.module_name:
-                affected_module = True
-            if affected_module:
-                all_affected_modules.append(module.module_name)
+        all_affected_modules = get_all_affected_modules_from_change(plain_module, plain_module_render_state)
 
         if len(all_affected_modules) > 0:
             choices[str(choice_idx)] = RenderChoice(
-                module=plain_module_render_state.change,
+                module=all_affected_modules[0],
                 render_range=None,
                 choice_type="rerender_affected",
                 wipe_later_modules=True,

tui/plain_module_render_choice_tui.py

@@ -1,11 +1,13 @@
+from functools import partial
+
 from textual import on
 from textual.app import App, ComposeResult
 from textual.binding import Binding
 from textual.containers import Vertical, VerticalScroll
 from textual.widgets import Label, ListItem, ListView, Static
 
 import plain_spec
-from partial_rendering import PlainModuleRenderState, RenderChoice
+from partial_rendering import PlainModuleRenderState, RenderChoice, get_all_affected_modules_from_change
 from plain_modules import PlainModule
 from tui.components import CustomFooter, TUIComponents
 
@@ -37,16 +39,12 @@ def __init__(
     def get_msg_from_choice(self, render_choice: RenderChoice) -> str:
         if render_choice.choice_type == "module_start":
             return f"Start from module [#5593FF]{render_choice.module.module_name}[/]"
-        elif render_choice.choice_type == "rerender_affected":
-            all_affected_modules = list[str]()
-            affected_module = False
-            for module in self.plain_module.all_required_modules:
-                if module.module_name == render_choice.module.module_name:
-                    affected_module = True
-                if affected_module:
-                    all_affected_modules.append(module.module_name)
-
-            return f"Re-render all affected modules ([#5593FF]{', '.join(all_affected_modules)}[/])"
+        elif render_choice.choice_type == "rerender_affected" and self.plain_module_render_state.change is not None:
+            all_affected_modules = get_all_affected_modules_from_change(
+                self.plain_module, self.plain_module_render_state
+            )
+
+            return f"Re-render all affected modules ([#5593FF]{', '.join([m.module_name for m in all_affected_modules])}[/])"
 
         elif render_choice.choice_type == "rerender_from_first":
             return f"Re-render from first module ([#5593FF]{render_choice.module.module_name}[/])"