Add command line LMFDB interface

Author: roed314

lmfdb_search

@@ -0,0 +1,321 @@
+#!/usr/bin/env -S sage -python
+
+# Other features we could add:
+# * Random object
+# * Raw SQL parsing
+# * Another database connection (other than devmirror)
+
+# TODO: Add json and tsv language support to downloads
+# TODO: Tests
+
+import sys
+import json
+from argparse import ArgumentParser
+from urllib.parse import parse_qs
+from lmfdb import db
+from lmfdb.utils.completeness import results_complete
+from psycodict.encoding import copy_dumps
+
+parser = ArgumentParser(
+    prog="LMFDBSearch",
+    description="""Command-line search interface for the L-functions and modular forms database (LMFDB)
+
+You can provide input in two formats:
+* A url string, matching the query-string part of an LMFDB url.  For example, searching for number fields with degree 4 would be done with degree=4, matching the LMFDB url https://www.lmfdb.org/NumberField/?degree=4.
+* A json dictionary such as {"degree":4}.
+""")
+
+parser.add_argument("table", help="Which table or section of the lmfdb to search")
+parser.add_argument("query", nargs="*", help="A query constraining which results are returned.  If not present, all results will be included")
+parser.add_argument("-i", "--input", help="Input file with search query", )
+parser.add_argument("-s", "--sql", action="store_true", help="Use SQL for input")
+parser.add_argument("-o", "--output", help="Output file for search results")
+parser.add_argument("-f", "--force", action="store_true", help="Overwrite output file even if it exists")
+parser.add_argument("-t", "--format", help="Output format", choices=["raw", "text", "json", "csv", "tsv", "sage", "gap", "pari", "magma", "oscar"], default="raw")
+parser.add_argument("-p", "--completeness", action="store_true", help="Whether to include completeness information at the begining of the output")
+parser.add_argument("-c", "--cols", help="Which columns to include in the output")
+parser.add_argument("-l", "--limit", type=int, help="The number of matching results to return (-1 for no limit)", default=50)
+parser.add_argument("--offset", type=int, help="Where to start in the list of results", default=0)
+parser.add_argument("--sort", help="columns by which to sort the search results; prepend a minus sign to a column to reverse it")
+parser.add_argument("--oneper", help="a list of columns, separated by commas.  If provided, only one result will be included with each given set of values for those columns (the first according to the provided sort order).")
+parser.add_argument("-d", "--debug", action="store_true", help="whether to raise errors")
+parser.add_argument("--host", help="PostgreSQL server host or socket directory [default: %(default)s]", default="devmirror.lmfdb.xyz")
+parser.add_argument("--port", type=int, help="PostreSQL server port [default: $(default)d]", default=5432)
+parser.add_argument("--user", help="PostgreSQL username [default: %(default)s]", default="lmfdb")
+parser.add_argument("--password", help="PostgreSQL password [default: %(default)s]", default="lmfdb")
+parser.add_argument("--dbname", help="PostgreSQL database name [default: %(default)s]", default="lmfdb")
+
+args = parser.parse_args()
+
+def section_lookup(section):
+    if section == "L":
+        from lmfdb.lfunctions.main import l_function_search, LFunctionSearchArray
+        return l_function_search, LFunctionSearchArray()
+    if section == "L/rational":
+        from lmfdb.lfunctions.main import l_function_search, LFunctionSearchArray
+        return l_function_search, LFunctionSearchArray(force_rational=True)
+    if section == "ModularForm/GL2/Q/holomorphic":
+        from lmfdb.classical_modular_forms.main import newform_search, CMFSearchArray
+        return newform_search, CMFSearchArray()
+    if section == "ModularForm/GL2/Q/Maass":
+        from lmfdb.maass.main import search, MaassSearchArray
+        return search, MaassSearchArray()
+    if section == "ModularForm/GL2/TotallyReal":
+        from lmfdb.hilbert_modular_forms.hilbert_modular_form import hilbert_modular_form_search, HMFSearchArray
+        return hilbert_modular_form_search, HMFSearchArray()
+    if section == "ModularForm/GL2/ImaginaryQuadratic":
+        from lmfdb.bianchi_modular_forms.bianchi_modular_form import bianchi_modular_form_search, BMFSearchArray
+        return bianchi_modular_form_search, BMFSearchArray()
+    if section == "EllipticCurve/Q":
+        from lmfdb.elliptic_curves.elliptic_curve import elliptic_curve_search, ECSearchArray
+        return elliptic_curve_search, ECSearchArray()
+    if section == "EllipticCurve":
+        from lmfdb.ecnf.main import elliptic_curve_search, ECNFSearchArray
+        return elliptic_curve_search, ECNFSearchArray()
+    if section == "Genus2Curve/Q":
+        from lmfdb.genus2_curves.main import genus2_curve_search, G2CSearchArray
+        return genus2_curve_search, G2CSearchArray()
+    if section == "ModularCurve/Q":
+        from lmfdb.modular_curves.main import modcurve_search, ModCurveSearchArray
+        return modcurve_search, ModCurveSearchArray()
+    if section == "HigherGenus/C/Aut":
+        from lmfdb.higher_genus_w_automorphisms.main import higher_genus_w_automorphisms_search, HGCWASearchArray
+        return higher_genus_w_automorphisms_search, HGCWASearchArray()
+    if section == "Variety/Abelian/Fq":
+        from lmfdb.abvar.fq.main import abelian_variety_search, AbvarSearchArray
+        return abelian_variety_search, AbvarSearchArray()
+    if section == "Belyi":
+        from lmfdb.belyi.main import belyi_search, BelyiSearchArray
+        return belyi_search, BelyiSearchArray()
+    if section == "NumberField":
+        from lmfdb.number_fields.number_field import number_field_search, NFSearchArray
+        return number_field_search, NFSearchArray()
+    if section == "padicField":
+        from lmfdb.local_fields.main import local_field_search, LFSearchArray
+        return local_field_search, LFSearchArray()
+    if section == "Character/Dirichlet/":
+        from lmfdb.characters.main import dirichlet_character_search, DirichSearchArray
+        return dirichlet_character_search, DirichSearchArray()
+    if section == "ArtinRepresentation":
+        from lmfdb.artin_representations.main import artin_representation_search, ArtinSearchArray
+        return artin_representation_search, ArtinSearchArray()
+    if section == "Motive/Hypergeometric/Q":
+        from lmfdb.hypergm.main import hgm_search, HGMSearchArray
+        return hgm_search, HGMSearchArray()
+    if section == "GaloisGroup":
+        from lmfdb.galois_groups.main import galois_group_search, GalSearchArray
+        return galois_group_search, GalSearchArray()
+    if section == "SatoTateGroup":
+        from lmfdb.sato_tate_groups import sato_tate_search, STSearchArray
+        return sato_tate_search, STSearchArray()
+    if section == "Groups/Abstract":
+        from lmfdb.groups.abstract.main import group_search, GroupsSearchArray
+        return group_search, GroupsSearchArray()
+    if section == "Lattice":
+        from lmfdb.lattice.main import lattice_search, LatSearchArray
+        return lattice_search, LatSearchArray()
+    # TODO: include tabs (e.g. families for p-adic fields, newspaces
+
+    parser.error("Must provide an LMFDB table or section as the first argument")
+
+def print_help(table, col=None):
+    # TODO: Add useful help
+    pass
+
+def write(s, F):
+    if F is None:
+        if s and s[-1] == "\n":
+            s = s[:-1]
+        print(s)
+    else:
+        _ = F.write(s)
+
+#if args.host or args.port or args.user or args.password or args.dbname:
+#    parser.error("Specifying a non default host is not yet supported")
+
+if args.cols is not None:
+    projection = args.cols = args.cols.split(",")
+    if any(col not in table.search_cols for col in projection):
+        parser.error("Invalid columns in output: " + ",".join(col for col in projection if col not in table.search_cols))
+else:
+    projection = 1
+
+if args.table == "help":
+    # Support giving a table name via a second positional argument, such as
+    if not args.query:
+        parser.print_help()
+    elif len(args.query) <= 2:
+        print_help(*args.query)
+    else:
+        parser.error("Too many positional arguments")
+    sys.exit(1)
+elif len(args.query) > 1:
+    parser.error("Too many positional arguments")
+elif args.table in db.tablenames:
+    section_search = False
+    table, search_array = db[args.table], None
+    if args.format not in ["raw", "json", "csv", "tsv"]:
+        parser.error(f"{args.format} output format not supported for LMFDB tables, only for sections")
+else:
+    section_search = True
+    wrapper, search_array = section_lookup(args.table)
+    table = wrapper.table
+    if ("download" not in wrapper.shortcuts and args.format not in ["raw", "json", "csv", "tsv"] or
+        "download" in wrapper.shortcuts and args.format not in ["raw"] + list(wrapper.shortcuts["download"].languages)):
+        parser.error(f"{args.format} output format not supported for {args.table}")
+    projection = [col for col in wrapper.projection if col in table.search_cols]
+
+if args.query and args.input is not None:
+    parser.error("may only provide input from one source: input file or positional argument")
+elif args.input is not None:
+    try:
+        with open(args.input) as F:
+            args.query = [F.read()]
+    except IOError as err:
+        if args.debug:
+            raise
+        parser.error(str(err))
+
+# Autodetect input format based on query
+# Must be compatible with table specification
+sort = one_per = raw = None
+if args.debug:
+    print("QUERY", args.query)
+if not args.query:
+    query = {}
+elif args.sql:
+    query = {}
+    raw = args.query[0]
+elif args.query[0][0] == "{":
+    try:
+        query = json.loads(args.query[0])
+    except Exception as err:
+        if args.debug:
+            raise
+        parser.error(str(err))
+else:
+    info = parse_qs(args.query[0])
+    for key, val in info.items():
+        if len(val) > 1:
+            parser.error(f"Multiple values found for {key} in query string")
+        info[key] = val[0]
+    info["search_array"] = search_array
+    query, sort, table, title, err_title, template, one_per = wrapper.make_query(info)
+
+if args.oneper:
+    # one_per might also have been set by url, but we allow the user to override with this setting
+    one_per = args.oneper.split(",")
+    if any(col not in table.search_cols for col in one_per):
+        parser.error(",".join(col for col in one_per if col not in table.search_cols) + " not columns of " + table.search_table)
+
+if args.sort:
+    # sort might also have been set by url, but we allow the user to override with this setting
+    sort = args.sort.split(",")
+    sort = [(col[1:], -1) if col and col[0] == "-" else (col, 1) for col in sort]
+    if any(col not in table.search_cols for col,asc in sort):
+        parser.error(",".join(col for col,asc in sort if col not in table.search_cols) + " not columns of " + table.search_table)
+
+if args.limit == -1:
+    args.limit = None
+
+res = table.search(query,
+                   projection=projection,
+                   limit=args.limit,
+                   offset=args.offset,
+                   sort=sort,
+                   one_per=one_per,
+                   raw=raw)
+if projection == 1:
+    projection = table.search_cols
+
+if section_search and args.format not in ["raw", "json"]:
+    try:
+        if wrapper.cleaners:
+            if limit is None:
+                res = list(res)
+            for v in res:
+                for name, func in wrapper.cleaners.items():
+                    v[name] = func(v)
+        if wrapper.postprocess is not None:
+            # This could be expensive, since postprocess functions assume that res usually has at most 50 items
+            res = wrapper.postprocess(res, info, query)
+    except Exception as err:
+        if args.debug:
+            raise
+        parser.error(str(err))
+if args.format == "json" and args.limit is None:
+    res = list(res)
+
+if args.output:
+    if os.path.exists(args.output) and not args.force:
+        parser.error("Output file already exists")
+    try:
+        fkwds = {}
+        if args.format in ["csv", "tsv"]:
+            fkwds["newline"] = ""
+        F = open(args.output, "w", **fkwds)
+    except Exception as err:
+        if args.debug:
+            raise
+        parser.error(str(err))
+elif args.format in ["csv", "tsv"]:
+    F = sys.stdout
+else:
+    F = None
+try:
+    if args.completeness:
+        complete, msg, caveat = results_complete(table.search_table, query, table._db, search_array)
+        if args.format == "json":
+            res = {
+                "complete": complete,
+                "completeness_msg": msg,
+                "completeness_caveat": caveat,
+                "results": res
+            }
+        elif complete:
+            write("Complete, since the LMFDB contains all " + msg + "\n", F)
+            if caveat is None:
+                write("No reliance on unproven conjectures\n", F)
+            else:
+                write("The completeness " + caveat + "\n", F)
+        else:
+            write("Not guaranteed complete\n\n", F)
+    if args.format == "raw":
+        write("|".join(projection) + "\n", F)
+        write("|".join([table.col_type[col] for col in projection]) + "\n\n", F)
+        for rec in res:
+            line = "|".join(copy_dumps(rec.get(col,None), table.col_type[col]) for col in projection)
+            write(line + "\n", F)
+    elif section_search and "download" in wrapper.shortcuts:
+        dl = wrapper.shortcuts["download"]
+        lang = dl.languages[args.format]
+        columns = wrapper.columns
+        if args.cols is None:
+            cols = [col for col in columns.columns_shown(info, rank=-1) if col.default(info)]
+        else:
+            cols = [col for col in columns.columns if col.name in args.cols]
+        for s in lang.assign_iter("data", lang.to_lang_iter(
+                map(
+                    lambda rec: [col.download(dl.postprocess(rec, info, query)) for col in cols],
+                    res))):
+            write(s, F)
+    elif args.format == "json":
+        write(json.dumps(res), F)
+    else:
+        if args.format == "csv":
+            delimiter = ","
+        elif args.format == "tsv":
+            delimiter="\t"
+        else:
+            raise ValueError(f"Format {args.format} not valid")
+        writer = csv.DictWriter(F, projection, delimiter=delimiter)
+        writer.writeheader()
+        for rec in res:
+            writer.writerow(rec)
+except Exception as err:
+    if args.debug:
+        raise
+    parser.error(str(err))
+finally:
+    if F is not None and F is not sys.stdout:
+        F.close()