Add search to client

Author: lsbardel

quantflow/cli/app.py

@@ -1,94 +1,32 @@
-import asyncio
 import os
 from dataclasses import dataclass, field
-from typing import Any, Self
+from typing import Any
 
 import click
-import dotenv
-import pandas as pd
-from asciichartpy import plot
-from ccy.cli.console import df_to_rich
 from prompt_toolkit import PromptSession
 from prompt_toolkit.history import FileHistory
 from rich.console import Console
 from rich.text import Text
 
-from quantflow.data.fmp import FMP
 
-from . import settings
-
-dotenv.load_dotenv()
-
-FREQUENCIES = tuple(FMP().historical_frequencies())
+from . import settings, commands
 
 
 @click.group()
 def qf() -> None:
     pass
 
 
-@qf.command()
-@click.argument("symbol")
-@click.pass_context
-def profile(ctx: click.Context, symbol: str) -> None:
-    """Company profile"""
-    app = QfApp.from_context(ctx)
-    data = asyncio.run(get_profile(symbol))[0]
-    app.print(data.pop("description"))
-    df = pd.DataFrame(data.items(), columns=["Key", "Value"])
-    app.print(df_to_rich(df))
-
-
-@qf.command()
-@click.argument("symbol")
-@click.option(
-    "-h",
-    "--height",
-    type=int,
-    default=20,
-    show_default=True,
-    help="Chart height",
-)
-@click.option(
-    "-l",
-    "--length",
-    type=int,
-    default=100,
-    show_default=True,
-    help="Number of data points",
-)
-@click.option(
-    "-f",
-    "--frequency",
-    type=click.Choice(FREQUENCIES),
-    default="",
-    help="Number of data points",
-)
-def chart(symbol: str, height: int, length: int, frequency: str) -> None:
-    """Symbol chart"""
-    df = asyncio.run(get_prices(symbol, frequency))
-    data = list(reversed(df["close"].tolist()[:length]))
-    print(plot(data, {"height": height}))
-
-
-async def get_prices(symbol: str, frequency: str) -> pd.DataFrame:
-    async with FMP() as cli:
-        return await cli.prices(symbol, frequency)
-
-
-async def get_profile(symbol: str) -> list[dict]:
-    async with FMP() as cli:
-        return await cli.profile(symbol)
+qf.add_command(commands.exit)
+qf.add_command(commands.profile)
+qf.add_command(commands.search)
+qf.add_command(commands.chart)
 
 
 @dataclass
 class QfApp:
     console: Console = field(default_factory=Console)
 
-    @classmethod
-    def from_context(cls, ctx: click.Context) -> Self:
-        return ctx.obj  # type: ignore
-
     def __call__(self) -> None:
         os.makedirs(settings.SETTINGS_DIRECTORY, exist_ok=True)
         history = FileHistory(str(settings.HIST_FILE_PATH))
@@ -123,12 +61,12 @@ def handle_command(self, text: str) -> None:
             return
         elif text == "help":
             return qf.main(["--help"], standalone_mode=False, obj=self)
-        elif text == "exit":
-            raise click.Abort()
 
         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:
             self.error(e)
+        except click.exceptions.UsageError as e:
+            self.error(e)

quantflow/cli/commands.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import click
+import asyncio
+import pandas as pd
+from typing import TYPE_CHECKING
+from asciichartpy import plot
+from ccy.cli.console import df_to_rich
+from quantflow.data.fmp import FMP
+
+FREQUENCIES = tuple(FMP().historical_frequencies())
+
+if TYPE_CHECKING:
+    from quantflow.cli.app import QfApp
+
+
+def from_context(ctx: click.Context) -> QfApp:
+    return ctx.obj  # type: ignore
+
+
+@click.command()
+def exit() -> None:
+    """Exit the program"""
+    raise click.Abort()
+
+
+@click.command()
+@click.argument("symbol")
+@click.pass_context
+def profile(ctx: click.Context, symbol: str) -> None:
+    """Company profile"""
+    app = from_context(ctx)
+    data = asyncio.run(get_profile(symbol))
+    if not data:
+        app.error(f"Company {symbol} not found - try searching")
+    else:
+        d = data[0]
+        app.print(d.pop("description") or "")
+        df = pd.DataFrame(d.items(), columns=["Key", "Value"])
+        app.print(df_to_rich(df))
+
+
+@click.command()
+@click.argument("text")
+@click.pass_context
+def search(ctx: click.Context, text: str) -> None:
+    """Search companies"""
+    app = from_context(ctx)
+    data = asyncio.run(search_company(text))
+    df = pd.DataFrame(data, columns=["symbol", "name", "currency", "stockExchange"])
+    app.print(df_to_rich(df))
+
+
+@click.command()
+@click.argument("symbol")
+@click.option(
+    "-h",
+    "--height",
+    type=int,
+    default=20,
+    show_default=True,
+    help="Chart height",
+)
+@click.option(
+    "-l",
+    "--length",
+    type=int,
+    default=100,
+    show_default=True,
+    help="Number of data points",
+)
+@click.option(
+    "-f",
+    "--frequency",
+    type=click.Choice(FREQUENCIES),
+    default="",
+    help="Frequency of data - if not provided it is daily",
+)
+def chart(symbol: str, height: int, length: int, frequency: str) -> None:
+    """Symbol chart"""
+    df = asyncio.run(get_prices(symbol, frequency))
+    if df.empty:
+        raise click.UsageError(
+            f"No data for {symbol} - are you sure the symbol exists?"
+        )
+    data = list(reversed(df["close"].tolist()[:length]))
+    print(plot(data, {"height": height}))
+
+
+async def get_prices(symbol: str, frequency: str) -> pd.DataFrame:
+    async with FMP() as cli:
+        return await cli.prices(symbol, frequency)
+
+
+async def get_profile(symbol: str) -> list[dict]:
+    async with FMP() as cli:
+        return await cli.profile(symbol)
+
+
+async def search_company(text: str) -> list[dict]:
+    async with FMP() as cli:
+        return await cli.search(text)

quantflow/cli/script.py

@@ -1,3 +1,7 @@
+import dotenv
+
+dotenv.load_dotenv()
+
 try:
     from .app import QfApp
 except ImportError: