From 3bf42b192c8f44136c34f6612be591a8e72a6714 Mon Sep 17 00:00:00 2001
From: Richard <gohn76@gmail.com>
Date: Wed, 17 Apr 2024 18:35:59 -0700
Subject: [PATCH] Update hot reloading logic in execute_module.py

This change separates the hot reloading logic for the mesop module and external app modules.

Before this fix, if the app module is outside of the mesop module, then changes to code in labs would not be reflected in the hot reload.
---
 mesop/cli/execute_module.py      | 16 ++++++++--------
 mesop/cli/execute_module_test.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/mesop/cli/execute_module.py b/mesop/cli/execute_module.py
index 15372aa3b..d516f0871 100644
--- a/mesop/cli/execute_module.py
+++ b/mesop/cli/execute_module.py
@@ -68,23 +68,23 @@ def get_app_modules(
 
   submodules: set[str] = set()
   for module in loaded_module_names:
+    module_segments = module.split(".")
+
     # Special case for mesop/example_index.py, which would normally consider
     # everything in `mesop` a sub-package/sub-module, however reloading modules
     # like runtime causes weird bugs. Thus, we only reload "app" modules and
     # not any core framework modules.
-    module_segments = module.split(".")
-    if (
-      module_segments[0] == "mesop"
-      and not (
+    if module_segments[0] == "mesop":
+      if (
         module_segments[:2] == ["mesop", "example_index"]
         or module_segments[:2] == ["mesop", "examples"]
         # Reset labs (b/c io.py needs to re-register stateclass)
         or module_segments[:2] == ["mesop", "labs"]
         or "e2e" in module_segments
-      )
-    ):
-      continue
-    if (
+      ):
+        submodules.add(module)
+    # We also want to hot reload user "app" modules.
+    elif (
       module.split(".")[: len(main_module_prefix_segments)]
       == main_module_prefix_segments
     ):
diff --git a/mesop/cli/execute_module_test.py b/mesop/cli/execute_module_test.py
index dd0a1cf79..88332440f 100644
--- a/mesop/cli/execute_module_test.py
+++ b/mesop/cli/execute_module_test.py
@@ -37,6 +37,37 @@ def test_get_app_modules():
   )
 
 
+def test_get_external_app_modules():
+  assert get_app_modules(
+    "external_app.app",
+    loaded_module_names=set(
+      [
+        "random",
+        "numpy",
+        "pandas",
+        "external_app.styles",
+        "external_app.components",
+        "mesop",
+        "mesop.runtime",
+        "mesop.example_index",
+        "mesop.examples",
+        "mesop.examples.index",
+        "mesop.components.radio.radio",
+        "mesop.components.radio.e2e.radio_app",
+      ]
+    ),
+  ) == set(
+    [
+      "external_app.styles",
+      "external_app.components",
+      "mesop.example_index",
+      "mesop.examples",
+      "mesop.examples.index",
+      "mesop.components.radio.e2e.radio_app",
+    ]
+  )
+
+
 if __name__ == "__main__":
   import pytest