@@ -861,8 +861,6 @@ def is_screen_installed(self, screen: Screen | str) -> bool:
861861 def get_screen (self , screen : Screen | str ) -> Screen :
862862 """Get an installed screen.
863863
864- If the screen isn't running, it will be registered before it is run.
865-
866864 Args:
867865 screen (Screen | str): Either a Screen object or screen name (the `name` argument when installed).
868866
@@ -879,10 +877,30 @@ def get_screen(self, screen: Screen | str) -> Screen:
879877 raise KeyError (f"No screen called { screen !r} installed" ) from None
880878 else :
881879 next_screen = screen
882- if not next_screen .is_running :
883- self ._register (self , next_screen )
884880 return next_screen
885881
882+ def _get_screen (self , screen : Screen | str ) -> tuple [Screen , AwaitMount ]:
883+ """Get an installed screen and a await mount object.
884+
885+ If the screen isn't running, it will be registered before it is run.
886+
887+ Args:
888+ screen (Screen | str): Either a Screen object or screen name (the `name` argument when installed).
889+
890+ Raises:
891+ KeyError: If the named screen doesn't exist.
892+
893+ Returns:
894+ tuple[Screen, AwaitMount]: A screen instance and an awaitable that awaits the children mounting.
895+
896+ """
897+ _screen = self .get_screen (screen )
898+ if not _screen .is_running :
899+ widgets = self ._register (self , _screen )
900+ return (_screen , AwaitMount (widgets ))
901+ else :
902+ return (_screen , AwaitMount ([]))
903+
886904 def _replace_screen (self , screen : Screen ) -> Screen :
887905 """Handle the replaced screen.
888906
@@ -900,19 +918,20 @@ def _replace_screen(self, screen: Screen) -> Screen:
900918 self .log .system (f"{ screen } REMOVED" )
901919 return screen
902920
903- def push_screen (self , screen : Screen | str ) -> None :
921+ def push_screen (self , screen : Screen | str ) -> AwaitMount :
904922 """Push a new screen on the screen stack.
905923
906924 Args:
907925 screen (Screen | str): A Screen instance or the name of an installed screen.
908926
909927 """
910- next_screen = self .get_screen (screen )
928+ next_screen , await_mount = self ._get_screen (screen )
911929 self ._screen_stack .append (next_screen )
912930 self .screen .post_message_no_wait (events .ScreenResume (self ))
913931 self .log .system (f"{ self .screen } is current (PUSHED)" )
932+ return await_mount
914933
915- def switch_screen (self , screen : Screen | str ) -> None :
934+ def switch_screen (self , screen : Screen | str ) -> AwaitMount :
916935 """Switch to another screen by replacing the top of the screen stack with a new screen.
917936
918937 Args:
@@ -921,12 +940,14 @@ def switch_screen(self, screen: Screen | str) -> None:
921940 """
922941 if self .screen is not screen :
923942 self ._replace_screen (self ._screen_stack .pop ())
924- next_screen = self .get_screen (screen )
943+ next_screen , await_mount = self ._get_screen (screen )
925944 self ._screen_stack .append (next_screen )
926945 self .screen .post_message_no_wait (events .ScreenResume (self ))
927946 self .log .system (f"{ self .screen } is current (SWITCHED)" )
947+ return await_mount
948+ return AwaitMount ([])
928949
929- def install_screen (self , screen : Screen , name : str | None = None ) -> str :
950+ def install_screen (self , screen : Screen , name : str | None = None ) -> AwaitMount :
930951 """Install a screen.
931952
932953 Args:
@@ -938,7 +959,7 @@ def install_screen(self, screen: Screen, name: str | None = None) -> str:
938959 ScreenError: If the screen can't be installed.
939960
940961 Returns:
941- str: The name of the screen
962+ AwaitMount: An awaitable that awaits the mounting of the screen and its children.
942963 """
943964 if name is None :
944965 name = nanoid .generate ()
@@ -949,9 +970,9 @@ def install_screen(self, screen: Screen, name: str | None = None) -> str:
949970 "Can't install screen; {screen!r} has already been installed"
950971 )
951972 self ._installed_screens [name ] = screen
952- self .get_screen (name ) # Ensures screen is running
973+ _screen , await_mount = self ._get_screen (name ) # Ensures screen is running
953974 self .log .system (f"{ screen } INSTALLED name={ name !r} " )
954- return name
975+ return await_mount
955976
956977 def uninstall_screen (self , screen : Screen | str ) -> str | None :
957978 """Uninstall a screen. If the screen was not previously installed then this
0 commit comments