Keep track of commands (#13)

Author: lsbardel

quantflow/cli/app.py

@@ -5,38 +5,23 @@
 
 import click
 from prompt_toolkit import PromptSession
+from prompt_toolkit.completion import NestedCompleter
 from prompt_toolkit.history import FileHistory
 from rich.console import Console
 from rich.text import Text
 
-from quantflow.data.fmp import FMP
-from quantflow.data.fred import Fred
 from quantflow.data.vault import Vault
 
 from . import settings
-from .commands import fred, stocks, vault
-
-
-@click.group()
-def qf() -> None:
-    pass
-
-
-@qf.command()
-def exit() -> None:
-    """Exit the program"""
-    raise click.Abort()
-
-
-qf.add_command(vault.vault)
-qf.add_command(stocks.stocks)
-qf.add_command(fred.fred)
+from .commands import quantflow
+from .commands.base import QuantGroup
 
 
 @dataclass
 class QfApp:
     console: Console = field(default_factory=Console)
     vault: Vault = field(default_factory=partial(Vault, settings.VAULT_FILE_PATH))
+    sections: list[QuantGroup] = field(default_factory=lambda: [quantflow])
 
     def __call__(self) -> None:
         os.makedirs(settings.SETTINGS_DIRECTORY, exist_ok=True)
@@ -49,14 +34,33 @@ def __call__(self) -> None:
         try:
             while True:
                 try:
-                    text = session.prompt("quantflow> ")
+                    text = session.prompt(
+                        self.prompt_message(),
+                        completer=self.prompt_completer(),
+                        complete_while_typing=True,
+                    )
                 except KeyboardInterrupt:
                     break
                 else:
                     self.handle_command(text)
         except click.Abort:
             self.console.print(Text("Bye!", style="bold magenta"))
 
+    def prompt_message(self) -> str:
+        name = ":".join([str(section.name) for section in self.sections])
+        return f"{name} > "
+
+    def prompt_completer(self) -> NestedCompleter:
+        return NestedCompleter.from_nested_dict(
+            {command: None for command in self.sections[-1].commands}
+        )
+
+    def set_section(self, section: QuantGroup) -> None:
+        self.sections.append(section)
+
+    def back(self) -> None:
+        self.sections.pop()
+
     def print(self, text_alike: Any, style: str = "") -> None:
         if isinstance(text_alike, str):
             style = style or "cyan"
@@ -67,29 +71,14 @@ def error(self, err: str | Exception) -> None:
         self.console.print(Text(f"\n{err}\n", style="bold red"))
 
     def handle_command(self, text: str) -> None:
-        self.current_command = text.split(" ")[0].strip()
         if not text:
             return
-        elif text == "help":
-            return qf.main(["--help"], standalone_mode=False, obj=self)
-
+        command = self.sections[-1]
         try:
-            qf.main(text.split(), standalone_mode=False, obj=self)
-        except click.exceptions.MissingParameter as e:
-            self.error(e)
-        except click.exceptions.NoSuchOption as e:
+            command.main(text.split(), standalone_mode=False, obj=self)
+        except (
+            click.exceptions.MissingParameter,
+            click.exceptions.NoSuchOption,
+            click.exceptions.UsageError,
+        ) as e:
             self.error(e)
-        except click.exceptions.UsageError as e:
-            self.error(e)
-
-    def fmp(self) -> FMP:
-        if key := self.vault.get("fmp"):
-            return FMP(key=key)
-        else:
-            raise click.UsageError("No FMP API key found")
-
-    def fred(self) -> Fred:
-        if key := self.vault.get("fred"):
-            return Fred(key=key)
-        else:
-            raise click.UsageError("No FRED API key found")

quantflow/cli/commands/__init__.py

@@ -0,0 +1,16 @@
+from .base import QuantContext, quant_group
+from .fred import fred
+from .stocks import stocks
+from .vault import vault
+
+
+@quant_group()
+def quantflow() -> None:
+    ctx = QuantContext.current()
+    if ctx.invoked_subcommand is None:
+        ctx.qf.print(ctx.get_help())
+
+
+quantflow.add_command(vault)
+quantflow.add_command(stocks)
+quantflow.add_command(fred)

quantflow/cli/commands/base.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Self, cast
+from typing import TYPE_CHECKING, Any, Self, cast
 
 import click
 
@@ -21,6 +21,12 @@ def current(cls) -> Self:
     def qf(self) -> QfApp:
         return self.obj  # type: ignore
 
+    def set_as_section(self) -> None:
+        group = cast(QuantGroup, self.command)
+        group.add_command(back)
+        self.qf.set_section(group)
+        self.qf.print(self.get_help())
+
     def fmp(self) -> FMP:
         if key := self.qf.vault.get("fmp"):
             return FMP(key=key)
@@ -41,3 +47,33 @@ class QuantCommand(click.Command):
 class QuantGroup(click.Group):
     context_class = QuantContext
     command_class = QuantCommand
+
+
+@click.command(cls=QuantCommand)
+def exit() -> None:
+    """Exit the program"""
+    raise click.Abort()
+
+
+@click.command(cls=QuantCommand)
+def help() -> None:
+    """display the commands"""
+    if ctx := QuantContext.current().parent:
+        cast(QuantContext, ctx).qf.print(ctx.get_help())
+
+
+@click.command(cls=QuantCommand)
+def back() -> None:
+    """Exit the current section"""
+    ctx = QuantContext.current()
+    ctx.qf.back()
+    ctx.qf.handle_command("help")
+
+
+def quant_group() -> Any:
+    return click.group(
+        cls=QuantGroup,
+        commands=[exit, help],
+        invoke_without_command=True,
+        add_help_option=False,
+    )

quantflow/cli/commands/fred.py

@@ -12,21 +12,20 @@
 
 from quantflow.data.fred import Fred
 
-from .base import QuantContext, QuantGroup
+from .base import QuantContext, quant_group
 
 FREQUENCIES = tuple(Fred.freq)
 
 if TYPE_CHECKING:
     pass
 
 
-@click.group(invoke_without_command=True, cls=QuantGroup)
+@quant_group()
 def fred() -> None:
-    """Federal Reserve of St. Louis data"""
+    """Federal Reserve of St. Louis data commands"""
     ctx = QuantContext.current()
     if ctx.invoked_subcommand is None:
-        ctx.qf.print("Welcome to FRED data commands!")
-        ctx.qf.print(ctx.get_help())
+        ctx.set_as_section()
 
 
 @fred.command()

quantflow/cli/commands/stocks.py

@@ -9,18 +9,17 @@
 
 from quantflow.data.fmp import FMP
 
-from .base import QuantContext, QuantGroup
+from .base import QuantContext, quant_group
 
 FREQUENCIES = tuple(FMP().historical_frequencies())
 
 
-@click.group(invoke_without_command=True, cls=QuantGroup)
+@quant_group()
 def stocks() -> None:
     """Stocks commands"""
     ctx = QuantContext.current()
     if ctx.invoked_subcommand is None:
-        ctx.qf.print("Welcome to the stocks commands!")
-        ctx.qf.print(ctx.get_help())
+        ctx.set_as_section()
 
 
 @stocks.command()

quantflow/cli/commands/vault.py

@@ -1,14 +1,16 @@
 import click
 
-from .base import QuantContext, QuantGroup
+from .base import QuantContext, quant_group
 
 API_KEYS = ("fmp", "fred")
 
 
-@click.group(invoke_without_command=True, cls=QuantGroup)
+@quant_group()
 def vault() -> None:
     """Manage vault secrets"""
-    pass
+    ctx = QuantContext.current()
+    if ctx.invoked_subcommand is None:
+        ctx.set_as_section()
 
 
 @vault.command()