Unity context managers

modal/_output.py

@@ -154,6 +154,7 @@ def finalize(self):
 
 class OutputManager:
     _instance: ClassVar[Optional["OutputManager"]] = None
+    _tree: ClassVar[Optional[Tree]] = None
 
     _console: Console
     _task_states: Dict[str, int]
@@ -167,7 +168,6 @@ class OutputManager:
     _app_page_url: Optional[str]
     _show_image_logs: bool
     _status_spinner_live: Optional[Live]
-    _tree: Optional[Tree]
 
     def __init__(
         self,
@@ -188,7 +188,6 @@ def __init__(
         self._app_page_url = None
         self._show_image_logs = False
         self._status_spinner_live = None
-        self._tree = None
 
     @classmethod
     def disable(cls):
@@ -387,17 +386,22 @@ def show_status_spinner(self):
         with self._status_spinner_live:
             yield
 
+    @classmethod
     @contextlib.contextmanager
-    def make_tree(self):
-        self._tree = Tree(step_progress("Creating objects..."), guide_style="gray50")
+    def make_tree(cls):
+        # Construct a tree even if the output isn't visible, but don't show it
+        cls._tree = Tree(step_progress("Creating objects..."), guide_style="gray50")
 
         try:
-            with self.make_live(self._tree):
+            if output_mgr := OutputManager.get():
+                with output_mgr.make_live(cls._tree):
+                    yield
+                cls._tree.label = step_completed("Created objects.")
+                output_mgr.print(output_mgr._tree)
+            else:
                 yield
-            self._tree.label = step_completed("Created objects.")
-            self.print(self._tree)
         finally:
-            self._tree = None
+            cls._tree = None
 
     def add_status_row(self) -> "StatusRow":
         return StatusRow(self._tree)

modal/runner.py

@@ -2,7 +2,6 @@
 import asyncio
 import dataclasses
 import os
-from contextlib import contextmanager
 from multiprocessing.synchronize import Event
 from typing import TYPE_CHECKING, Any, AsyncGenerator, Coroutine, Dict, List, Optional, TypeVar
 
@@ -101,15 +100,6 @@ async def _init_local_app_from_name(
         )
 
 
-@contextmanager
-def display():
-    if output_mgr := OutputManager.get():
-        with output_mgr.make_tree():
-            yield
-    else:
-        yield
-
-
 async def _create_all_objects(
     client: _Client,
     running_app: RunningApp,
@@ -126,7 +116,7 @@ async def _create_all_objects(
         environment_name=environment_name,
         app_id=running_app.app_id,
     )
-    with display():
+    with OutputManager.make_tree():
         # Get current objects, and reset all objects
         tag_to_object_id = running_app.tag_to_object_id
         running_app.tag_to_object_id = {}