Add Colab support with colab_run

Author: wwwillchen

mesop/__init__.py

@@ -107,6 +107,7 @@
 from mesop.features import page as page
 from mesop.key import Key as Key
 from mesop.runtime import runtime
+from mesop.server.colab_run import colab_run as colab_run
 
 _T = TypeVar("_T")
 

mesop/cli/cli.py

@@ -14,7 +14,9 @@
   reset_runtime,
   runtime,
 )
+from mesop.server.constants import EDITOR_PACKAGE_PATH, PROD_PACKAGE_PATH
 from mesop.server.flags import port
+from mesop.server.logging import log_startup
 from mesop.server.server import configure_flask_app
 from mesop.server.static_file_serving import configure_static_file_serving
 
@@ -71,23 +73,16 @@ def main(argv):
 
   configure_static_file_serving(
     flask_app,
-    static_file_runfiles_base="mesop/mesop/web/src/app/prod/web_package"
+    static_file_runfiles_base=PROD_PACKAGE_PATH
     if FLAGS.prod
-    else "mesop/mesop/web/src/app/editor/web_package",
+    else EDITOR_PACKAGE_PATH,
     livereload_script_url=os.environ.get("IBAZEL_LIVERELOAD_URL"),
   )
 
   if FLAGS.verbose:
     logging.getLogger("werkzeug").setLevel(logging.INFO)
   else:
-    # Log basic information about where the server is running since
-    # we don't get it printed by werkzeurg:
-    # https://github.com/pallets/werkzeug/blob/eafbed0ce2a6bdf60e62de82bf4a8365188ac334/src/werkzeug/serving.py#L820C9-L820C17
-    FgGreen = "\x1b[32m"
-    Reset = "\x1b[0m"
-    print(
-      f"\n{FgGreen}Running server on: http://localhost:{port()}{Reset}",
-    )
+    log_startup()
     logging.getLogger("werkzeug").setLevel(logging.WARN)
 
   flask_app.run(host="0.0.0.0", port=port(), use_reloader=False)

mesop/server/colab_run.py

@@ -0,0 +1,29 @@
+import threading
+
+from mesop.runtime import enable_debug_mode
+from mesop.server.constants import EDITOR_PACKAGE_PATH
+from mesop.server.flags import port
+from mesop.server.logging import log_startup
+from mesop.server.server import configure_flask_app
+from mesop.server.static_file_serving import configure_static_file_serving
+
+
+def colab_run():
+  """
+  Typically only used for Colab/Jupyter notebooks, otherwise you'll use the CLI
+  to execute a Mesop application.
+  """
+  flask_app = configure_flask_app()
+  enable_debug_mode()
+  configure_static_file_serving(
+    flask_app, static_file_runfiles_base=EDITOR_PACKAGE_PATH
+  )
+
+  log_startup()
+
+  def run_flask_app():
+    flask_app.run(host="0.0.0.0", port=port(), use_reloader=False)
+
+  # Launch Flask in background thread so we don't hog up the main thread
+  # for regular Colab usage.
+  threading.Thread(target=run_flask_app).start()

mesop/server/constants.py

@@ -0,0 +1,2 @@
+EDITOR_PACKAGE_PATH = "mesop/mesop/web/src/app/editor/web_package"
+PROD_PACKAGE_PATH = "mesop/mesop/web/src/app/prod/web_package"

mesop/server/logging.py

@@ -0,0 +1,12 @@
+from mesop.server.flags import port
+
+
+def log_startup():
+  # Log basic information about where the server is running since
+  # we don't get it printed by werkzeurg:
+  # https://github.com/pallets/werkzeug/blob/eafbed0ce2a6bdf60e62de82bf4a8365188ac334/src/werkzeug/serving.py#L820C9-L820C17
+  FgGreen = "\x1b[32m"
+  Reset = "\x1b[0m"
+  print(
+    f"\n{FgGreen}Running server on: http://localhost:{port()}{Reset}",
+  )

mesop/utils/caller.py

@@ -3,7 +3,9 @@
 import mesop.protos.ui_pb2 as pb
 
 
-def get_caller_source_code_location(levels: int = 1) -> pb.SourceCodeLocation:
+def get_caller_source_code_location(
+  levels: int = 1,
+) -> pb.SourceCodeLocation | None:
   current_frame = inspect.currentframe()
 
   # Walk backwards
@@ -20,7 +22,11 @@ def get_caller_source_code_location(levels: int = 1) -> pb.SourceCodeLocation:
   caller_info = inspect.getframeinfo(current_frame)
 
   # Get module from filepath
-  segments = caller_info.filename.split("runfiles/")[1].split("/")
+  splitted_file = caller_info.filename.split("runfiles/")
+  # Can't always look up from runfiles (e.g. colab).
+  if len(splitted_file) < 2:
+    return None
+  segments = splitted_file[1].split("/")
   segments[len(segments) - 1] = segments[len(segments) - 1][: len(".py") * -1]
   module = ".".join(segments)