From e32e094b9265d1690a6e4df3e1964161ed69d70f Mon Sep 17 00:00:00 2001
From: darrenburns <darrenburns@users.noreply.github.com>
Date: Wed, 16 Nov 2022 15:47:48 +0000
Subject: [PATCH] Support callables in App.SCREENS (#1185)

* Support Type[Screen] in App.SCREENS (lazy screens)

* Update CHANGELOG

* Remove redundant isinstance
---
 CHANGELOG.md          |  1 +
 src/textual/app.py    | 12 ++++++++----
 tests/test_screens.py | 31 ++++++++++++++++++++++++++++++-
 3 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c300711d7..7148331ba 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
   https://github.com/Textualize/textual/issues/1094
 - Added Pilot.wait_for_animation
 - Added `Widget.move_child` https://github.com/Textualize/textual/issues/1121
+- Support lazy-instantiated Screens (callables in App.SCREENS) https://github.com/Textualize/textual/pull/1185
 
 ### Changed
 
diff --git a/src/textual/app.py b/src/textual/app.py
index b5526f73c..e74a90f8a 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -25,6 +25,7 @@ from typing import (
     TypeVar,
     Union,
     cast,
+    Callable,
 )
 from weakref import WeakSet, WeakValueDictionary
 
@@ -228,7 +229,7 @@ class App(Generic[ReturnType], DOMNode):
     }
     """
 
-    SCREENS: dict[str, Screen] = {}
+    SCREENS: dict[str, Screen | Callable[[], Screen]] = {}
     _BASE_PATH: str | None = None
     CSS_PATH: CSSPathType = None
     TITLE: str | None = None
@@ -330,7 +331,7 @@ class App(Generic[ReturnType], DOMNode):
         self._registry: WeakSet[DOMNode] = WeakSet()
 
         self._installed_screens: WeakValueDictionary[
-            str, Screen
+            str, Screen | Callable[[], Screen]
         ] = WeakValueDictionary()
         self._installed_screens.update(**self.SCREENS)
 
@@ -998,12 +999,15 @@ class App(Generic[ReturnType], DOMNode):
                 next_screen = self._installed_screens[screen]
             except KeyError:
                 raise KeyError(f"No screen called {screen!r} installed") from None
+            if callable(next_screen):
+                next_screen = next_screen()
+                self._installed_screens[screen] = next_screen
         else:
             next_screen = screen
         return next_screen
 
     def _get_screen(self, screen: Screen | str) -> tuple[Screen, AwaitMount]:
-        """Get an installed screen and a await mount object.
+        """Get an installed screen and an AwaitMount object.
 
         If the screen isn't running, it will be registered before it is run.
 
@@ -1558,7 +1562,7 @@ class App(Generic[ReturnType], DOMNode):
 
         # Close pre-defined screens
         for screen in self.SCREENS.values():
-            if screen._running:
+            if isinstance(screen, Screen) and screen._running:
                 await self._prune_node(screen)
 
         # Close any remaining nodes
diff --git a/tests/test_screens.py b/tests/test_screens.py
index 0841faf51..707bad5df 100644
--- a/tests/test_screens.py
+++ b/tests/test_screens.py
@@ -11,8 +11,37 @@ skip_py310 = pytest.mark.skipif(
 )
 
 
+async def test_installed_screens():
+    class ScreensApp(App):
+        SCREENS = {
+            "home": Screen,  # Screen type
+            "one": Screen(),  # Screen instance
+            "two": lambda: Screen()  # Callable[[], Screen]
+        }
+
+    app = ScreensApp()
+    async with app.run_test() as pilot:
+        pilot.app.push_screen("home")  # Instantiates and pushes the "home" screen
+        pilot.app.push_screen("one")   # Pushes the pre-instantiated "one" screen
+        pilot.app.push_screen("home")  # Pushes the single instance of "home" screen
+        pilot.app.push_screen("two")   # Calls the callable, pushes returned Screen instance
+
+        assert len(app.screen_stack) == 5
+        assert app.screen_stack[1] is app.screen_stack[3]
+        assert app.screen is app.screen_stack[4]
+        assert isinstance(app.screen, Screen)
+        assert app.is_screen_installed(app.screen)
+
+        assert pilot.app.pop_screen()
+        assert pilot.app.pop_screen()
+        assert pilot.app.pop_screen()
+        assert pilot.app.pop_screen()
+        with pytest.raises(ScreenStackError):
+            pilot.app.pop_screen()
+
+
+
 @skip_py310
-@pytest.mark.asyncio
 async def test_screens():
 
     app = App()