From efe0342a6f040ca135b74c5414181cb57608a01b Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Tue, 16 Aug 2022 20:39:43 +0100
Subject: [PATCH] docstring

---
 src/textual/app.py  | 92 +++++++++++++++++++++++++--------------------
 tests/test_focus.py |  2 +
 2 files changed, 54 insertions(+), 40 deletions(-)

diff --git a/src/textual/app.py b/src/textual/app.py
index ca432e77e..d5ef2331d 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -639,51 +639,18 @@ class App(Generic[ReturnType], DOMNode):
             self._register(self, next_screen)
         return next_screen
 
-    def push_screen(self, screen: Screen | str) -> None:
-        """Push a new screen on the screen stack.
+    def _replace_screen(self, screen: Screen, remove: bool | None = None) -> Screen:
+        """Handle the replaced screen.
 
         Args:
-            screen (Screen | str): A Screen instance or an id.
+            screen (Screen): A screen object.
+            remove (bool | None): Remove replaced screen if True. Don't remove if False.
+                If None, remove screens not in self.SCREENS.
 
         Returns:
-            Screen: Newly active screen.
+            Screen: The replaced screen
         """
-        next_screen = self._get_screen(screen)
-        self._screen_stack.append(next_screen)
-        self.screen.post_message_no_wait(events.ScreenResume(self))
-
-    def switch_screen(self, screen: Screen | str) -> Screen:
-        """Switch to a another screen.
-
-        Args:
-            screen (Screen | str): A screen instance or a named of a screen.
-
-        Returns:
-            Screen: The previous screen object.
-        """
-        next_screen = self._get_screen(screen)
-        current_screen = self._screen_stack.pop()
-        current_screen.post_message_no_wait(events.ScreenSuspend(self))
-        self._screen_stack.append(next_screen)
-        self.screen.post_message_no_wait(events.ScreenResume(self))
-        return current_screen
-
-    def pop_screen(self, remove: bool | None = None) -> Screen:
-        """Pop the current screen from the stack, and switch to the previous screen.
-
-        Returns:
-            Screen: The screen that was replaced.
-        """
-        screen_stack = self._screen_stack
-        if len(screen_stack) <= 1:
-            raise ScreenStackError(
-                "Can't pop screen; there must be at least one screen on the stack"
-            )
-        screen = screen_stack.pop()
         screen.post_message_no_wait(events.ScreenSuspend(self))
-        self.screen._screen_resized(self.size)
-        self.screen.post_message_no_wait(events.ScreenResume(self))
-
         if remove is None:
             if screen not in self.SCREENS.values():
                 screen.remove()
@@ -696,9 +663,54 @@ class App(Generic[ReturnType], DOMNode):
                 screen.remove()
             else:
                 screen.detach()
-
         return screen
 
+    def push_screen(self, screen: Screen | str) -> None:
+        """Push a new screen on the screen stack.
+
+        Args:
+            screen (Screen | str): A Screen instance or an id.
+
+        """
+        next_screen = self._get_screen(screen)
+        self._screen_stack.append(next_screen)
+        self.screen.post_message_no_wait(events.ScreenResume(self))
+
+    def switch_screen(self, screen: Screen | str, remove: bool | None) -> Screen:
+        """Switch to a another screen.
+
+        Args:
+            screen (Screen | str): A screen instance or a named of a screen.
+            remove (bool | None): Remove replaced screen if True. Don't remove if False.
+                If None, remove screens not in self.SCREENS.
+
+        Returns:
+            Screen: The previous screen object.
+        """
+        next_screen = self._get_screen(screen)
+        previous_screen = self._replace_screen(self._screen_stack.pop(), remove=remove)
+        self._screen_stack.append(next_screen)
+        self.screen.post_message_no_wait(events.ScreenResume(self))
+        return previous_screen
+
+    def pop_screen(self, remove: bool | None = None) -> Screen:
+        """Pop the current screen from the stack, and switch to the previous screen.
+
+        Returns:
+            Screen: The screen that was replaced.
+            remove (bool | None): Remove replaced screen if True. Don't remove if False.
+                If None, remove screens not in self.SCREENS.
+        """
+        screen_stack = self._screen_stack
+        if len(screen_stack) <= 1:
+            raise ScreenStackError(
+                "Can't pop screen; there must be at least one screen on the stack"
+            )
+        previous_screen = self._replace_screen(screen_stack.pop(), remove)
+        self.screen._screen_resized(self.size)
+        self.screen.post_message_no_wait(events.ScreenResume(self))
+        return previous_screen
+
     def set_focus(self, widget: Widget | None) -> None:
         """Focus (or unfocus) a widget. A focused widget will receive key events first.
 
diff --git a/tests/test_focus.py b/tests/test_focus.py
index dfe8dad8a..8d62ab901 100644
--- a/tests/test_focus.py
+++ b/tests/test_focus.py
@@ -14,6 +14,7 @@ class NonFocusable(Widget, can_focus=False, can_focus_children=False):
 async def test_focus_chain():
 
     app = App()
+    app._set_active()
     app.push_screen(Screen())
 
     # Check empty focus chain
@@ -34,6 +35,7 @@ async def test_focus_chain():
 async def test_focus_next_and_previous():
 
     app = App()
+    app._set_active()
     app.push_screen(Screen())
     app.screen.add_children(
         Focusable(id="foo"),