mantis-imf 0.1.8__py3-none-any.whl

mantis_imf/__init__.py

@@ -0,0 +1,293 @@
+from __future__ import annotations
+
+import argparse
+import json
+import shlex
+import sys
+from collections.abc import Mapping, Sequence
+from datetime import datetime
+from typing import Any
+
+from asof_store import AsOfStore
+
+from .ib_models import Instrument
+
+_SEARCH_FIELDS = ("symbol", "local_symbol", "country", "exchange", "product_type")
+_STORAGE_FIELD_ALIASES = {
+    "localSymbol": "local_symbol",
+    "exchange": "exchange_id",
+    "exchangeId": "exchange_id",
+    "productType": "product_type",
+    "type": "product_type",
+}
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="imf",
+        epilog=(
+            "Scraper commands: scrape-instruments, scrape-reconcile, "
+            "scrape-funds, scrape-exchanges"
+        ),
+    )
+    commands = parser.add_subparsers(dest="command", required=True)
+    search = commands.add_parser("search", help="search stored instruments")
+    _add_store_arguments(search)
+    for field in _SEARCH_FIELDS:
+        search.add_argument(f"--{field.replace('_', '-')}", help=f"match {field}")
+    search.add_argument(
+        "--field",
+        action="append",
+        default=[],
+        metavar="NAME=VALUE",
+        help="match any JSON field; dotted names address nested fields (repeatable)",
+    )
+    search.add_argument(
+        "--exact", action="store_true", help="match values exactly, ignoring case"
+    )
+    index = commands.add_parser(
+        "index", help="create PostgreSQL indexes for fast fuzzy searches"
+    )
+    _add_store_arguments(index)
+    fuzzy = commands.add_parser(
+        "fuzzy", help="search instruments for every supplied value across fields"
+    )
+    _add_store_arguments(fuzzy)
+    fuzzy.add_argument("values", nargs="*", metavar="VALUE")
+    fuzzy.add_argument(
+        "--exact", action="store_true", help="match values exactly, ignoring case"
+    )
+    return parser
+
+
+def _add_store_arguments(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--sql-uri",
+        help="override the SQL URI from service-registrar.toml",
+    )
+    parser.add_argument("--table-name", default="ib_instruments")
+    parser.add_argument(
+        "--missing-primary-key-table-name", default="ib_instrument_templates"
+    )
+    parser.add_argument(
+        "--as-of",
+        type=datetime.fromisoformat,
+        help="only include instrument versions at or before this ISO timestamp",
+    )
+
+
+def _json_value(value: Any) -> Any:
+    if isinstance(value, Instrument):
+        return value.model_dump(mode="json", by_alias=True, exclude_unset=True)
+    if isinstance(value, Mapping):
+        return value
+    return None
+
+
+def _filters(args: argparse.Namespace) -> list[tuple[str, str]]:
+    filters: list[tuple[str, str]] = []
+    for field in _SEARCH_FIELDS:
+        query = getattr(args, field)
+        if query is not None:
+            # asof-store persists Instrument models using Python field names.
+            storage_field = {
+                "local_symbol": "local_symbol",
+                "exchange": "exchange_id",
+                "product_type": "product_type",
+            }.get(field, field)
+            filters.append((storage_field, query))
+    for expression in args.field:
+        name, separator, query = expression.partition("=")
+        if not separator or not name or not query:
+            raise ValueError(f"invalid --field {expression!r}; expected NAME=VALUE")
+        parts = name.split(".")
+        parts[0] = _STORAGE_FIELD_ALIASES.get(parts[0], parts[0])
+        filters.append((".".join(parts), query))
+    return filters
+
+
+def search_instruments(argv: Sequence[str] | None = None) -> int:
+    args = _build_parser().parse_args(argv)
+    if args.command != "search":
+        raise ValueError(f"unsupported command: {args.command}")
+    return _search(args)
+
+
+def fuzzy_instruments(argv: Sequence[str] | None = None) -> int:
+    args = _build_parser().parse_args(argv)
+    if args.command != "fuzzy":
+        raise ValueError(f"unsupported command: {args.command}")
+    if not args.values:
+        return _fuzzy_repl(args)
+    return _search(args)
+
+
+def _search(args: argparse.Namespace) -> int:
+    _validate_table_names(args)
+
+    store = AsOfStore.from_sql(
+        sql_uri=args.sql_uri,
+        table_name=args.table_name,
+        timestamp_type=datetime,
+        key_type=int,
+        value_type=Instrument,
+    )
+    template_store = None
+    try:
+        template_store = AsOfStore.from_sql(
+            sql_uri=args.sql_uri,
+            table_name=args.missing_primary_key_table_name,
+            timestamp_type=datetime,
+            key_type=Instrument,
+            value_type=type(None),
+        )
+        _search_records(args, store, template_store)
+    finally:
+        store.close()
+        if template_store is not None:
+            template_store.close()
+    return 0
+
+
+def _validate_table_names(args: argparse.Namespace) -> None:
+    if args.table_name == args.missing_primary_key_table_name:
+        raise ValueError("instrument and missing-primary-key table names must differ")
+
+
+def _create_trigram_indexes(store: Any, column: str, fields: tuple[str, ...]) -> None:
+    create_indexes = getattr(store, "create_trigram_indexes", None)
+    if create_indexes is None:
+        raise ValueError(f"{type(store).__name__} does not support trigram indexes")
+    create_indexes(fields, column=column)
+
+
+def _create_search_indexes(args: argparse.Namespace) -> int:
+    _validate_table_names(args)
+    fields = tuple(Instrument.model_fields)
+    store = AsOfStore.from_sql(
+        sql_uri=args.sql_uri,
+        table_name=args.table_name,
+        timestamp_type=datetime,
+        key_type=int,
+        value_type=Instrument,
+    )
+    template_store = None
+    try:
+        template_store = AsOfStore.from_sql(
+            sql_uri=args.sql_uri,
+            table_name=args.missing_primary_key_table_name,
+            timestamp_type=datetime,
+            key_type=Instrument,
+            value_type=type(None),
+        )
+        _create_trigram_indexes(store, "value", fields)
+        _create_trigram_indexes(template_store, "key", fields)
+        print(
+            "Ensured PostgreSQL trigram search indexes for "
+            f"{len(fields)} instrument fields."
+        )
+    finally:
+        store.close()
+        if template_store is not None:
+            template_store.close()
+    return 0
+
+
+def _search_records(args: argparse.Namespace, store: Any, template_store: Any) -> None:
+    if args.command == "fuzzy":
+        fields = tuple(Instrument.model_fields)
+        terms = tuple(args.values)
+        records = store.as_of(args.as_of).get_all_matching_terms(
+            terms, fields, exact=args.exact
+        )
+        templates = template_store.as_of(args.as_of).get_keys_matching_terms(
+            terms, fields, exact=args.exact
+        )
+    else:
+        filters = _filters(args)
+        if filters:
+            records = store.as_of(args.as_of).get_all_matching(
+                filters, exact=args.exact
+            )
+            templates = template_store.as_of(args.as_of).get_keys_matching(
+                filters, exact=args.exact
+            )
+        else:
+            records = store.get_all(args.as_of)
+            templates = template_store.get_keys_matching([], as_of=args.as_of)
+    for source_records in (records, templates):
+        for value in source_records:
+            record = _json_value(value)
+            if record is not None:
+                print(json.dumps(record, ensure_ascii=False))
+
+
+def _fuzzy_repl(args: argparse.Namespace) -> int:
+    if args.table_name == args.missing_primary_key_table_name:
+        raise ValueError("instrument and missing-primary-key table names must differ")
+
+    store = AsOfStore.from_sql(
+        sql_uri=args.sql_uri,
+        table_name=args.table_name,
+        timestamp_type=datetime,
+        key_type=int,
+        value_type=Instrument,
+    )
+    template_store = None
+    try:
+        template_store = AsOfStore.from_sql(
+            sql_uri=args.sql_uri,
+            table_name=args.missing_primary_key_table_name,
+            timestamp_type=datetime,
+            key_type=Instrument,
+            value_type=type(None),
+        )
+        print("Enter values to search; use exit or quit to leave.")
+        while True:
+            try:
+                line = input("imf fuzzy> ")
+            except EOFError:
+                print()
+                break
+            except KeyboardInterrupt:
+                print()
+                break
+            try:
+                terms = shlex.split(line)
+            except ValueError as exc:
+                print(f"imf fuzzy: {exc}", file=sys.stderr)
+                continue
+            if not terms:
+                continue
+            if len(terms) == 1 and terms[0].casefold() in {"exit", "quit"}:
+                break
+            args.values = terms
+            _search_records(args, store, template_store)
+    finally:
+        store.close()
+        if template_store is not None:
+            template_store.close()
+    return 0
+
+
+def main() -> None:
+    try:
+        argv = sys.argv[1:]
+        if argv and argv[0].startswith("scrape-"):
+            from .scraper import main as scraper_main
+
+            raise SystemExit(scraper_main(argv))
+        args = _build_parser().parse_args(argv)
+        if args.command == "index":
+            result = _create_search_indexes(args)
+        elif args.command == "fuzzy" and not args.values:
+            result = _fuzzy_repl(args)
+        else:
+            result = _search(args)
+        raise SystemExit(result)
+    except (ValueError, ImportError) as exc:
+        print(f"imf: error: {exc}", file=sys.stderr)
+        raise SystemExit(2) from exc
+
+
+__all__ = ["fuzzy_instruments", "main", "search_instruments"]