From eaa749665f9b8271eff45be8e5e1e72ac8729b9e Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Thu, 14 Sep 2023 10:26:51 +0100
Subject: [PATCH] change hotkey

---
 .../guide/command_palette/command01.py        |  2 +-
 docs/guide/command_palette.md                 | 13 ++++++---
 src/textual/command.py                        | 27 ++++++++++++++++---
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/docs/examples/guide/command_palette/command01.py b/docs/examples/guide/command_palette/command01.py
index 0efa25a120..43a4ea860a 100644
--- a/docs/examples/guide/command_palette/command01.py
+++ b/docs/examples/guide/command_palette/command01.py
@@ -15,7 +15,7 @@ def read_files(self) -> list[Path]:
         """Get a list of Python files in the current working directory."""
         return list(Path("./").glob("*.py"))
 
-    async def post_init(self) -> None:  # (1)!
+    async def startup(self) -> None:  # (1)!
         """Called once when the command palette is opened, prior to searching."""
         worker = self.app.run_worker(self.read_files, thread=True)
         self.python_paths = await worker.wait()
diff --git a/docs/guide/command_palette.md b/docs/guide/command_palette.md
index e3e36dab4b..214e7f5850 100644
--- a/docs/guide/command_palette.md
+++ b/docs/guide/command_palette.md
@@ -68,13 +68,13 @@ The following example will display a blank screen initially, but if you hit ++ct
   5. Highlights matching letters in the search.
   6. Adds our custom command source and the default command sources.
 
-There are two methods you will typically override in a command source: [`post_init`][textual.command.Source.post_init] and [`search`][textual.command.Source.search].
-Both should be coroutines (`async def`).
+There are three methods you can override in a command source: [`startup`][textual.command.Source.startup], [`search`][textual.command.Source.search], and [`shutdown`][textual.command.Source.shutdown].
+All of these methods should be coroutines (`async def`). Only `search` is required, the other methods are optional.
 Let's explore those methods in detail.
 
-### post_init method
+### startup method
 
-The [`post_init`][textual.command.Source.post_init] method is called when the command palette is opened.
+The [`startup`][textual.command.Source.startup] method is called when the command palette is opened.
 You can use this method as way of performing work that needs to be done prior to searching.
 In the example, we use this method to get the Python (.py) files in the current working directory.
 
@@ -99,6 +99,11 @@ In the example above, the callback is a lambda which calls the `open_file` metho
     This is a deliberate design decision taken to prevent a single broken `Source` class from making the command palette unusable.
     Errors in command sources will be logged to the [console](./devtools.md).
 
+### Shutdown method
+
+The [`shutdown`][textual.command.Source.shutdown] method is called when the command palette is closed.
+You can use this as a hook to gracefully close any objects you created in startup.
+
 ## Screen commands
 
 You can also associate commands with a screen by adding a `COMMAND_SOURCES` class var to your Screen class.
diff --git a/src/textual/command.py b/src/textual/command.py
index e263b3ff8d..66c6304810 100644
--- a/src/textual/command.py
+++ b/src/textual/command.py
@@ -7,12 +7,13 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from asyncio import CancelledError, Queue, Task, TimeoutError, wait_for
+from asyncio import Queue, Task, wait, wait_for
 from dataclasses import dataclass
 from functools import total_ordering
 from time import monotonic
-from typing import TYPE_CHECKING, Any, AsyncGenerator, AsyncIterator, ClassVar
+from typing import TYPE_CHECKING, Any, AsyncIterator, ClassVar
 
+import rich.repr
 from rich.align import Align
 from rich.console import Group, RenderableType
 from rich.emoji import Emoji
@@ -162,7 +163,7 @@ def _post_init(self) -> None:
         async def post_init_task() -> None:
             """Wrapper to post init that runs in a task."""
             try:
-                await self.post_init()
+                await self.startup()
             except Exception:
                 self.app.log.error(Traceback())
             else:
@@ -175,7 +176,7 @@ async def _wait_init(self) -> None:
         if self._init_task is not None:
             await self._init_task
 
-    async def post_init(self) -> None:
+    async def startup(self) -> None:
         """Called after the Source is initialized, but before any calls to `search`."""
 
     async def _search(self, query: str) -> Hits:
@@ -205,7 +206,16 @@ async def search(self, query: str) -> Hits:
         """
         yield NotImplemented
 
+    async def shutdown(self) -> None:
+        """Called when the Source is shutdown.
 
+        Use this method to perform an cleanup, if required.
+
+        """
+        print("SHUTDOWN Source ")
+
+
+@rich.repr.auto
 @total_ordering
 class Command(Option):
     """Class that holds a command in the [`CommandList`][textual.command.CommandList]."""
@@ -510,6 +520,15 @@ def on_mount(self, _: Mount) -> None:
         for _source in self._sources:
             _source._post_init()
 
+    async def on_unmount(self) -> None:
+        try:
+            await wait(
+                [create_task(source.shutdown()) for source in self._sources],
+            )
+        except Exception:
+            self.log(Exception)
+        self._sources.clear()
+
     def _stop_busy_countdown(self) -> None:
         """Stop any busy countdown that's in effect."""
         if self._busy_timer is not None: