esgpull 0.6.3__py3-none-any.whl

esgpull/cli/convert.py

@@ -0,0 +1,285 @@
+from collections import Counter
+from collections.abc import MutableMapping
+from pathlib import Path
+
+import click
+import pyparsing as pp
+import yaml
+from click.exceptions import Abort, Exit
+from rich.box import MINIMAL_DOUBLE_HEAD
+from rich.prompt import Confirm
+from rich.table import Table
+from rich.text import Text
+
+from esgpull.cli.decorators import opts
+from esgpull.cli.utils import init_esgpull
+from esgpull.graph import Graph
+from esgpull.models import Options, Query, Tag
+from esgpull.models.selection import FacetValues, Selection
+from esgpull.tui import Verbosity, logger
+
+SKIP = {"priority", "protocol"}
+options_names = Options()._names
+pp.ParserElement.set_default_whitespace_chars(" \t,")  # remove newline
+
+
+def line(expr: pp.ParserElement) -> pp.ParserElement:
+    return pp.LineStart().suppress() + expr + pp.LineEnd().suppress()
+
+
+def sqbr(expr: pp.ParserElement) -> pp.ParserElement:
+    return pp.Suppress("[") + expr + pp.Suppress("]")
+
+
+anything = pp.Word(pp.printables)
+word = pp.Word(pp.alphanums + "-_")
+eq = pp.Suppress("=")
+
+comment_start = pp.Char("#")[1, ...].suppress()
+name_comment = pp.Group(line(comment_start + anything[1])[0, 1])
+comment = line(comment_start + anything[...]).suppress()
+facet_name = word("name")
+facet_value = word
+facet_values = pp.Group(word[1, ...])
+facet = pp.Group(line(facet_name + eq + facet_values("vals")))("facet")
+variable_cmip5 = pp.Group(
+    line(
+        "variable"
+        + sqbr(facet_value)("realm")
+        + sqbr(facet_value)("time_frequency")
+        + eq
+        + facet_values("variable")
+    )
+)("variable_cmip5")
+variable_cmip6 = pp.Group(
+    line(
+        "variable"
+        + sqbr(pp.Opt("table_id=") + facet_value("table_id"))
+        + eq
+        + facet_values("variable_id")
+    )
+)("variable_cmip6")
+variable_no_sqbr = pp.Group(line("variable" + eq + facet_values("variable")))(
+    "variable_no_sqbr"
+)
+otherwise = line(anything[...])("otherwise")
+rest = pp.Group(
+    comment
+    | facet
+    | variable_no_sqbr
+    | variable_cmip5
+    | variable_cmip6
+    | otherwise
+)[...]
+selection_file = name_comment("name") + rest("rest")
+
+
+def remove_duplicates(
+    selection: MutableMapping[str, FacetValues],
+) -> MutableMapping[str, FacetValues]:
+    result: dict[str, FacetValues] = {}
+    duplicates: dict[str, list[str]] = {}
+    for name, values in selection.items():
+        if isinstance(values, str):
+            result[name] = values
+            continue
+        counter = Counter(values)
+        nb_dup = sum(c > 1 for c in counter.values())
+        if nb_dup == 0:
+            result[name] = values
+            continue
+        duplicates[name] = list(dict(counter.most_common(nb_dup)))
+    if duplicates:
+        logger.warning(f"Duplicate values {duplicates}'")
+    return result
+
+
+def is_CMIP6(q: Query) -> bool:
+    project = set(q.selection["project"] + q.selection["mip_era"])
+    return "CMIP6" in project
+
+
+def isnot_CMIP6(q: Query) -> bool:
+    project = set(q.selection["project"] + q.selection["mip_era"])
+    return any({"CMIP5", "CORDEX"} & project)
+
+
+def fix_CMIP5(q: Query) -> Query:
+    qd = q.asdict()
+    if q.selection["frequency"]:
+        qd["selection"]["time_frequency"] = qd["selection"].pop("frequency")
+    return Query(**qd)
+
+
+def convert_file(path: Path) -> Graph:
+    logger.info(path)
+    query = Query()
+    kids: list[Query] = []
+    result = selection_file.parse_file(path)
+    if result.name:
+        # name = result.name[0].split("@", 1)[0]
+        query.tags.append(Tag(name=result.name[0]))
+    for line in result.rest:
+        if line.facet:
+            name, values = line.facet.as_list()
+            if name in SKIP:
+                continue
+            elif name in options_names:
+                if len(values) > 1:
+                    raise ValueError({name: values})
+                d = {name: values[0]}
+                logger.debug(f"OPTION {d}")
+                setattr(query.options, name, values[0])
+            elif name in Selection._facet_names:
+                d = {name: values}
+                logger.debug(f"FACET {d}")
+                query.selection[name] = list(set(values))
+            else:
+                raise ValueError(f"{name!r} undefined\n{path.read_text()}")
+        elif line.variable_cmip5:
+            selection = line.variable_cmip5.as_dict()
+            logger.debug(f"SUBQUERY {selection}")
+            kid = Query(selection=remove_duplicates(selection), tracked=True)
+            kids.append(kid)
+        elif line.variable_cmip6:
+            selection = line.variable_cmip6.as_dict()
+            project = query.selection["project"]
+            if project and project[0] in ["CMIP5", "CORDEX"]:
+                if "variable_id" in selection:
+                    selection["variable"] = selection.pop("variable_id")
+                if "table_id" in selection:
+                    selection["time_frequency"] = selection.pop("table_id")
+            logger.debug(f"SUBQUERY {selection}")
+            kid = Query(selection=remove_duplicates(selection))
+            kids.append(kid)
+        elif line.variable_no_sqbr:
+            logger.error(line.as_dict())
+        elif line.as_list():
+            logger.error(line.as_list())
+    if len(kids) > 1:
+        query.untrack()
+    elif len(kids) == 1:
+        query = query << kids.pop()
+    else:
+        query.track(query.options)
+    if isnot_CMIP6(query):
+        query = fix_CMIP5(query)
+    query.compute_sha()
+    graph = Graph(None)
+    graph.add(query)
+    for kid in kids:
+        if isnot_CMIP6(query) or isnot_CMIP6(kid):
+            kid = fix_CMIP5(kid)
+        kid.require = query.sha
+        kid.compute_sha()
+        expanded = query << kid
+        kid.track(expanded.options)
+        graph.add(kid)
+    return graph
+
+
+CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"])
+
+
+@click.command(context_settings=CONTEXT_SETTINGS)
+@click.argument(
+    "paths",
+    required=True,
+    type=click.Path(exists=True, path_type=Path),
+    nargs=-1,
+)
+@click.option(
+    "--out",
+    "-o",
+    type=click.Path(exists=False, path_type=Path),
+    default=None,
+)
+@click.option("print_table", "--table", is_flag=True, default=False)
+@click.option("print_graph", "--graph", is_flag=True, default=False)
+@opts.record
+@opts.verbosity
+def convert(
+    paths: list[Path],
+    out: Path | None,
+    print_table: bool,
+    print_graph: bool,
+    record: bool,
+    verbosity: Verbosity,
+):
+    """
+    Convert synda selection files to esgpull queries
+
+    The flags `--table/--graph` are used to select the console output design.
+
+    Use `--out <path/to/output.yaml>` to generate a query file containing the converted synda
+    selection files. This file can be used as an input to the `add --query-file` command.
+
+    Note that `convert` takes any number of input paths, and produces a single output.
+    Using `**/*` as the input path is a good way to convert a whole tree of selection files.
+
+    As a an arbitrary convention, the generated queries will be tagged with the content of the first line of a synda selection file, if that line starts with `#` and has a single word (+symbols) without whitespaces in it.
+    """
+    esg = init_esgpull(verbosity, safe=False, record=record)
+    with esg.ui.logging("convert", onraise=Abort):
+        if len(paths) == 0:
+            esg.ui.print("No file provided")
+            raise click.exceptions.Exit(0)
+        if out is not None:
+            out_str = f"[yellow]{out.resolve()}[/]"
+            if out.is_file():
+                esg.ui.print(f"Overwriting existing file {out_str}")
+                allow = Confirm.ask("Continue?", default=False)
+                if not allow:
+                    raise click.exceptions.Exit(0)
+        if out is None and not print_table and not print_graph:
+            esg.ui.print(
+                "[red]Error[/]: At least one of "
+                "[yellow]--out/--table/--graph[/] is required."
+            )
+            raise click.exceptions.Exit(1)
+        table = Table(
+            box=MINIMAL_DOUBLE_HEAD,
+            show_edge=False,
+            show_lines=True,
+        )
+        table.add_column(Text("path", justify="left"))
+        table.add_column(Text("file", justify="center"))
+        table.add_column(Text("query", justify="center"))
+        full_graph = Graph(None)
+        empty_query = Query()
+        empty_query.options.apply_defaults(empty_query.options)
+        empty_query.compute_sha()
+        for path in paths:
+            if path.is_file():
+                try:
+                    graph = convert_file(path)
+                except Exception as e:
+                    msg = f"Cannot convert {path.resolve()}"
+                    logger.exception(e)
+                    logger.error(msg)
+                    esg.ui.print(f"[red]Error[/]: {msg}")
+                    continue
+                queries = list(graph.queries.values())
+                if len(queries) == 1 and queries[0].sha == empty_query.sha:
+                    esg.ui.print(
+                        f"Skipping empty query produced from {path.resolve()}"
+                    )
+                    continue
+                path_text = Text(str(path).replace("/", "/\n"), style="yellow")
+                file_text = Text(path.read_text())
+                table.add_row(path_text, file_text, graph)
+                full_graph.add(*queries, force=True)
+        if print_table:
+            esg.ui.print(table)
+        if print_graph:
+            esg.ui.print(full_graph)
+        if out is not None:
+            graph_as_yaml = yaml.dump(full_graph.dump())
+            try:
+                with out.open("w") as f:
+                    f.write(graph_as_yaml)
+                esg.ui.print(f":+1: [green]Graph was written to[/] {out_str}")
+            except Exception:
+                out.unlink()
+                raise
+        esg.ui.raise_maybe_record(Exit(0))

esgpull/cli/decorators.py

@@ -0,0 +1,342 @@
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, TypeAlias, TypeVar
+
+import click
+from click_params import StringListParamType
+
+from esgpull.cli.utils import EnumParam
+from esgpull.models import FileStatus, Option
+from esgpull.tui import Verbosity
+
+F = TypeVar("F", bound=Callable[..., Any])
+Dec: TypeAlias = Callable[[F], F]
+
+
+def compose(*decs: Dec) -> Dec:
+    def deco(fn: Callable):
+        for dec in reversed(decs):
+            fn = dec(fn)
+        return fn
+
+    return deco
+
+
+class args:
+    facets: Dec = click.argument(
+        "facets",
+        type=str,
+        nargs=-1,
+    )
+    key: Dec = click.argument(
+        "key",
+        type=str,
+        nargs=1,
+        required=False,
+        default=None,
+    )
+    value: Dec = click.argument(
+        "value",
+        type=str,
+        nargs=1,
+        required=False,
+        default=None,
+    )
+    path: Dec = click.argument(
+        "path",
+        type=click.Path(exists=False, path_type=Path),
+        required=False,
+        default=None,
+    )
+    status: Dec = click.argument(
+        "status",
+        type=EnumParam(FileStatus),
+        nargs=-1,
+    )
+    query_id: Dec = click.argument(
+        "query_id",
+        type=str,
+        nargs=1,
+        required=False,
+    )
+    query_ids: Dec = click.argument(
+        "query_ids",
+        type=str,
+        nargs=-1,
+        required=True,
+    )
+
+
+class opts:
+    all: Dec = click.option(
+        "all_",
+        "--all",
+        "-a",
+        is_flag=True,
+        default=False,
+    )
+    children: Dec = click.option(
+        "--children",
+        "-c",
+        is_flag=True,
+        default=False,
+    )
+    default: Dec = click.option(
+        "--default",
+        "-d",
+        is_flag=True,
+        default=False,
+    )
+    detail: Dec = click.option(
+        "--detail",
+        type=int,
+        default=None,
+    )
+    disable_ssl: Dec = click.option(
+        "--disable-ssl",
+        is_flag=True,
+        default=False,
+    )
+    dry_run: Dec = click.option(
+        "--dry-run",
+        "-z",
+        is_flag=True,
+        default=False,
+    )
+    facets_hints: Dec = click.option(
+        "facets_hints",
+        "--facets",
+        "-F",
+        is_flag=True,
+        default=False,
+    )
+    file: Dec = click.option(
+        "--file",
+        "-f",
+        is_flag=True,
+        default=False,
+    )
+    files: Dec = click.option(
+        "--files",
+        is_flag=True,
+        default=False,
+    )
+    force: Dec = click.option(
+        "--force",
+        is_flag=True,
+        default=False,
+    )
+    generate: Dec = click.option(
+        "--generate",
+        is_flag=True,
+        default=False,
+    )
+    hints: Dec = click.option(
+        "--hints",
+        "-H",
+        type=StringListParamType(","),
+        default=None,
+    )
+    name: Dec = click.option(
+        "--name",
+        "-n",
+        type=str,
+        default=None,
+    )
+    query_file: Dec = click.option(
+        "--query-file",
+        "-q",
+        type=click.Path(exists=True, dir_okay=False, path_type=Path),
+        default=None,
+    )
+    quiet: Dec = click.option(
+        "--quiet",
+        is_flag=True,
+        default=False,
+    )
+    record: Dec = click.option(
+        "--record",
+        is_flag=True,
+        default=False,
+    )
+    reset: Dec = click.option(
+        "--reset",
+        is_flag=True,
+        default=False,
+    )
+    shas: Dec = click.option(
+        "--shas",
+        "-s",
+        is_flag=True,
+        default=False,
+    )
+    show: Dec = click.option(
+        "--show",
+        is_flag=True,
+        default=False,
+    )
+    simple: Dec = click.option(
+        "--simple",
+        is_flag=True,
+        default=False,
+    )
+    since: Dec = click.option(
+        "--since",
+        type=str,
+        default=None,
+    )
+    status: Dec = click.option(
+        "--status",
+        type=EnumParam(FileStatus),
+        default=None,
+        multiple=True,
+    )
+    tag: Dec = click.option(
+        "--tag",
+        "-t",
+        type=str,
+        default=None,
+    )
+    track: Dec = click.option(
+        "--track",
+        is_flag=True,
+        default=False,
+    )
+    yes: Dec = click.option(
+        "--yes",
+        "-y",
+        is_flag=True,
+        default=False,
+    )
+    verbosity: Dec = click.option(
+        "verbosity",
+        "-v",
+        count=True,
+        type=EnumParam(Verbosity),
+    )
+
+
+# Display group
+_page: Dec = click.option(
+    "--page",
+    "-p",
+    type=int,
+    default=0,  # 1?
+)
+_zero: Dec = click.option(
+    "--zero",
+    "-0",
+    is_flag=True,
+    default=False,
+)
+_all: Dec = click.option(
+    "_all",
+    "--all",
+    "-a",
+    is_flag=True,
+    default=False,
+)
+
+# Json/Yaml exclusive group
+_json: Dec = click.option(
+    "--json",
+    is_flag=True,
+    default=False,
+)
+_yaml: Dec = click.option(
+    "--yaml",
+    is_flag=True,
+    default=False,
+)
+
+# Query definition group
+OptionChoice = click.Choice([opt.name for opt in list(Option)[:-1]])
+_tag: Dec = click.option("tags", "--tag", "-t", multiple=True)
+_require: Dec = click.option(
+    "--require",
+    "-r",
+    type=str,
+    nargs=1,
+    required=False,
+    default=None,
+)
+_distrib: Dec = click.option(
+    "--distrib",
+    "-d",
+    type=OptionChoice,
+)
+_latest: Dec = click.option(
+    "--latest",
+    type=OptionChoice,
+)
+_replica: Dec = click.option(
+    "--replica",
+    type=OptionChoice,
+)
+_retracted: Dec = click.option(
+    "--retracted",
+    type=OptionChoice,
+)
+
+# Query dates group
+datetime_type = click.DateTime(["%Y-%m-%d"])
+_from: Dec = click.option(
+    "date_from",
+    "--from",
+    type=datetime_type,
+    default=None,
+)
+_to: Dec = click.option(
+    "date_to",
+    "--to",
+    type=datetime_type,
+    default=None,
+)
+
+_children: Dec = click.option(
+    "--children",
+    "-c",
+    is_flag=True,
+    default=False,
+)
+_parents: Dec = click.option(
+    "--parents",
+    "-p",
+    is_flag=True,
+    default=False,
+)
+_expand: Dec = click.option(
+    "--expand",
+    "-e",
+    is_flag=True,
+    default=False,
+)
+
+
+class groups:
+    display: Dec = compose(
+        _page,
+        _zero,
+        _all,
+    )
+    json_yaml: Dec = compose(
+        _json,
+        _yaml,
+    )
+    query_def: Dec = compose(
+        _tag,
+        _require,
+        _distrib,
+        _latest,
+        _replica,
+        _retracted,
+    )
+    query_date: Dec = compose(
+        _from,
+        _to,
+    )
+    show: Dec = compose(
+        _children,
+        _parents,
+        _expand,
+    )

esgpull/cli/download.py

@@ -0,0 +1,74 @@
+import asyncio
+import sys
+
+import click
+import rich
+from click.exceptions import Abort, Exit
+
+if sys.version_info < (3, 11):
+    from exceptiongroup import BaseExceptionGroup
+
+from esgpull.cli.decorators import args, opts
+from esgpull.cli.utils import get_queries, init_esgpull, valid_name_tag
+from esgpull.models import File, FileStatus
+from esgpull.tui import Verbosity, logger
+from esgpull.utils import format_size
+
+
+@click.command()
+@args.query_id
+@opts.tag
+@opts.disable_ssl
+@opts.quiet
+@opts.record
+@opts.verbosity
+def download(
+    query_id: str | None,
+    tag: str | None,
+    disable_ssl: bool,
+    quiet: bool,
+    record: bool,
+    verbosity: Verbosity,
+):
+    """
+    Asynchronously download files linked to queries
+    """
+    esg = init_esgpull(verbosity, record=record)
+    if disable_ssl:
+        esg.config.download.disable_ssl = True
+    with esg.ui.logging("download", onraise=Abort):
+        if not valid_name_tag(esg.graph, esg.ui, query_id, tag):
+            esg.ui.raise_maybe_record(Exit(1))
+        if query_id is None and tag is None:
+            esg.graph.load_db()
+            graph = esg.graph
+        else:
+            queries = get_queries(esg.graph, query_id, tag)
+            graph = esg.graph.subgraph(
+                *queries,
+                children=True,
+                parents=True,
+            )
+        esg.ui.print(graph)
+        shas: set[str] = set()
+        queue: list[File] = []
+        for query in graph.queries.values():
+            for file in query.files:
+                if file.status == FileStatus.Queued and file.sha not in shas:
+                    shas.add(file.sha)
+                    queue.append(file)
+        if not queue:
+            rich.print("Download queue is empty.")
+            esg.ui.raise_maybe_record(Exit(0))
+        coro = esg.download(queue, show_progress=not quiet)
+        files, errors = asyncio.run(coro)
+        if files:
+            size = format_size(sum(file.size for file in files))
+            esg.ui.print(
+                f"Downloaded {len(files)} new files for a total size of {size}"
+            )
+        if errors:
+            logger.error(f"{len(errors)} files could not be installed.")
+            exc_group = BaseExceptionGroup("Download", [e.err for e in errors])
+            esg.ui.raise_maybe_record(exc_group)
+        esg.ui.raise_maybe_record(Exit(0))