esgpull 0.6.3__py3-none-any.whl

esgpull/cli/utils.py

@@ -0,0 +1,247 @@
+import sys
+from collections import OrderedDict
+from enum import Enum
+from pathlib import Path
+from typing import Any, Literal, MutableMapping, Sequence
+
+import click
+import yaml
+from click.exceptions import BadArgumentUsage
+from rich.box import MINIMAL_DOUBLE_HEAD
+from rich.table import Table
+from rich.text import Text
+
+from esgpull import Esgpull
+from esgpull.graph import Graph
+from esgpull.models import Dataset, File, Option, Options, Query, Selection
+from esgpull.tui import UI, TempUI, Verbosity
+from esgpull.utils import format_size
+
+
+def get_command() -> Text:
+    exe, *args = sys.argv
+    args = [arg for arg in args if arg != "--record"]
+    return Text(" ".join(["$", Path(exe).name, *args]))
+
+
+def init_esgpull(
+    verbosity: Verbosity,
+    safe: bool = True,
+    record: bool = False,
+    load_db: bool = True,
+) -> Esgpull:
+    TempUI.verbosity = Verbosity.Errors
+    with TempUI.logging():
+        esg = Esgpull(
+            verbosity=verbosity,
+            safe=safe,
+            record=record,
+            load_db=load_db,
+        )
+        if record:
+            esg.ui.print(get_command())
+    return esg
+
+
+class Messages:
+    @staticmethod
+    def no_such_query(name: str) -> str:
+        return f":stop_sign: No such query [green]{name}[/]"
+
+    @staticmethod
+    def none_tagged(tag: str) -> str:
+        return f":stop_sign: No query tagged with [magenta]{tag}[/]"
+
+    @staticmethod
+    def multimatch(name: str) -> str:
+        return f"Found multiple queries starting with [b cyan]{name}[/]"
+
+
+class EnumParam(click.Choice):
+    name = "enum"
+
+    def __init__(self, enum: type[Enum]):
+        self.__enum = enum
+        super().__init__(choices=[item.value for item in enum])
+
+    def convert(self, value, param, ctx) -> Enum:
+        converted_str = super().convert(value, param, ctx)
+        return self.__enum(converted_str)
+
+
+def filter_keys(
+    docs: Sequence[File | Dataset],
+    ids: range,
+    size: bool = True,
+    data_node: bool = False,
+    # date: bool = False,
+) -> list[OrderedDict[str, Any]]:
+    result: list[OrderedDict[str, Any]] = []
+    for i, doc in zip(ids, docs):
+        od: OrderedDict[str, Any] = OrderedDict()
+        od["id"] = str(i)
+        if isinstance(doc, File):
+            od["file"] = doc.file_id
+        else:
+            od["dataset"] = doc.dataset_id
+            od["#"] = str(doc.number_of_files)
+        if size:
+            od["size"] = doc.size
+        if data_node:
+            od["data_node"] = doc.data_node
+        # if date:
+        #     od["date"] = doc.get("timestamp") or doc.get("_timestamp")
+        result.append(od)
+    return result
+
+
+def totable(docs: list[OrderedDict[str, Any]]) -> Table:
+    table = Table(box=MINIMAL_DOUBLE_HEAD, show_edge=False)
+    for key in docs[0].keys():
+        justify: Literal["left", "right", "center"]
+        if key in ["file", "dataset"]:
+            justify = "left"
+        else:
+            justify = "right"
+        table.add_column(
+            Text(key, justify="center"),
+            justify=justify,
+        )
+    for doc in docs:
+        row: list[str] = []
+        for key, value in doc.items():
+            if key == "size":
+                value = format_size(value)
+            row.append(value)
+        table.add_row(*row)
+    return table
+
+
+def parse_facets(facets: list[str]) -> Selection:
+    facet_dict: dict[str, list[str]] = {}
+    exact_terms: list[str] | None = None
+    for facet in facets:
+        match facet.split(":"):
+            case [value]:
+                name = "query"
+            case [name, value] if name and value:
+                ...
+            case _:
+                raise BadArgumentUsage(f"{facet!r} is not valid syntax.")
+        if value.startswith("/"):
+            if exact_terms is not None:
+                raise BadArgumentUsage("Nested exact string is forbidden.")
+            exact_terms = []
+        if exact_terms is not None:
+            if name != "query":
+                raise BadArgumentUsage(
+                    "Cannot use facet term inside an exact string."
+                )
+            exact_terms.append(value)
+            if value.endswith("/"):
+                final_exact_str = " ".join(exact_terms)
+                value = '"' + final_exact_str.strip("/") + '"'
+                exact_terms = None
+            else:
+                continue
+            facet_dict.setdefault(name, [])
+            facet_dict[name].append(value)
+        else:
+            values = value.split(",")
+            facet_dict.setdefault(name, [])
+            facet_dict[name].extend(values)
+    selection = Selection()
+    for name, values in facet_dict.items():
+        selection[name] = values
+    return selection
+
+
+def parse_query(
+    facets: list[str],
+    # query options
+    tags: list[str],
+    require: str | None,
+    distrib: str | None,
+    latest: str | None,
+    replica: str | None,
+    retracted: str | None,
+) -> Query:
+    options = Options(
+        distrib=distrib or Option.notset,
+        latest=latest or Option.notset,
+        replica=replica or Option.notset,
+        retracted=retracted or Option.notset,
+    )
+    selection = parse_facets(facets)
+    return Query(
+        tags=tags,
+        require=require,
+        options=options,
+        selection=selection,
+    )
+
+
+def is_list_of_maps(seq: Sequence) -> bool:
+    return all(isinstance(item, MutableMapping) for item in seq)
+
+
+def serialize_queries_from_file(path: Path) -> list[Query]:
+    with path.open() as f:
+        content = yaml.safe_load(f)
+    queries: list[MutableMapping]
+    if isinstance(content, list):
+        queries = content
+    elif isinstance(content, MutableMapping):
+        values = list(content.values())
+        if is_list_of_maps(values):
+            queries = values
+        else:
+            queries = [content]
+    else:
+        raise ValueError(content)
+    return [Query(**query) for query in queries]
+
+
+def valid_name_tag(
+    graph: Graph,
+    ui: UI,
+    query_id: str | None,
+    tag: str | None,
+) -> bool:
+    result = True
+    if query_id is not None:
+        shas = graph.matching_shas(query_id, graph._shas)
+        if len(shas) > 1:
+            ui.print(Messages.multimatch(query_id))
+            ui.print(shas, json=True)
+            result = False
+        elif len(shas) == 0:
+            ui.print(Messages.no_such_query(query_id), err=True)
+            result = False
+    elif tag is not None:
+        tags = [t.name for t in graph.get_tags()]
+        if tag not in tags:
+            ui.print(Messages.none_tagged(tag), err=True)
+            result = False
+    return result
+
+
+def get_queries(
+    graph: Graph,
+    query_id: str | None,
+    tag: str | None,
+    children: bool = False,
+) -> list[Query]:
+    queries: list[Query] = []
+    if query_id is not None:
+        query = graph.get(name=query_id)
+        if query is not None:
+            queries = [query]
+    elif tag is not None:
+        queries = graph.with_tag(tag)
+    if children:
+        for i in range(len(queries)):
+            query = queries[0]
+            kids = graph.get_all_children(query.sha)
+            queries.extend(kids)
+    return queries

esgpull/config.py

@@ -0,0 +1,410 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Container, Iterator, Mapping
+from enum import Enum, auto
+from pathlib import Path
+from typing import Any, cast
+
+import tomlkit
+from attrs import Factory, define, field
+from cattrs import Converter
+from cattrs.gen import make_dict_unstructure_fn, override
+from tomlkit import TOMLDocument
+
+from esgpull.constants import CONFIG_FILENAME
+from esgpull.exceptions import BadConfigError, VirtualConfigError
+from esgpull.install_config import InstallConfig
+from esgpull.models.options import Options
+
+logger = logging.getLogger("esgpull")
+
+
+@define
+class Paths:
+    auth: Path = field(converter=Path)
+    data: Path = field(converter=Path)
+    db: Path = field(converter=Path)
+    log: Path = field(converter=Path)
+    tmp: Path = field(converter=Path)
+
+    @auth.default
+    def _auth_factory(self) -> Path:
+        if InstallConfig.current is not None:
+            root = InstallConfig.current.path
+        else:
+            root = InstallConfig.default
+        return root / "auth"
+
+    @data.default
+    def _data_factory(self) -> Path:
+        if InstallConfig.current is not None:
+            root = InstallConfig.current.path
+        else:
+            root = InstallConfig.default
+        return root / "data"
+
+    @db.default
+    def _db_factory(self) -> Path:
+        if InstallConfig.current is not None:
+            root = InstallConfig.current.path
+        else:
+            root = InstallConfig.default
+        return root / "db"
+
+    @log.default
+    def _log_factory(self) -> Path:
+        if InstallConfig.current is not None:
+            root = InstallConfig.current.path
+        else:
+            root = InstallConfig.default
+        return root / "log"
+
+    @tmp.default
+    def _tmp_factory(self) -> Path:
+        if InstallConfig.current is not None:
+            root = InstallConfig.current.path
+        else:
+            root = InstallConfig.default
+        return root / "tmp"
+
+    def __iter__(self) -> Iterator[Path]:
+        yield self.auth
+        yield self.data
+        yield self.db
+        yield self.log
+        yield self.tmp
+
+
+@define
+class Credentials:
+    filename: str = "credentials.toml"
+
+
+@define
+class Cli:
+    page_size: int = 20
+
+
+@define
+class Db:
+    filename: str = "esgpull.db"
+
+
+@define
+class Download:
+    chunk_size: int = 1 << 26  # 64 MiB
+    http_timeout: int = 20
+    max_concurrent: int = 5
+    disable_ssl: bool = False
+    disable_checksum: bool = False
+
+
+@define
+class DefaultOptions:
+    distrib: str = Options._distrib_.name
+    latest: str = Options._latest_.name
+    replica: str = Options._replica_.name
+    retracted: str = Options._retracted_.name
+
+    def asdict(self) -> dict[str, str]:
+        return dict(
+            distrib=self.distrib,
+            latest=self.latest,
+            replica=self.replica,
+            retracted=self.retracted,
+        )
+
+
+@define
+class API:
+    index_node: str = "esgf-node.ipsl.upmc.fr"
+    http_timeout: int = 20
+    max_concurrent: int = 5
+    page_limit: int = 50
+    default_options: DefaultOptions = Factory(DefaultOptions)
+
+
+def fix_rename_search_api(doc: TOMLDocument) -> TOMLDocument:
+    if "api" in doc and "search" in doc:
+        raise KeyError(
+            "Both 'api' and 'search' (deprecated) are used in your "
+            "config, please use 'api' only."
+        )
+    elif "search" in doc:
+        logger.warn(
+            "Deprecated key 'search' is used in your config, "
+            "please use 'api' instead."
+        )
+        doc["api"] = doc.pop("search")
+    return doc
+
+
+config_fixers = [fix_rename_search_api]
+
+
+class ConfigKind(Enum):
+    Virtual = auto()
+    NoFile = auto()
+    Partial = auto()
+    Complete = auto()
+
+
+class ConfigKey:
+    path: tuple[str, ...]
+
+    def __init__(self, first: str | tuple[str, ...], *rest: str) -> None:
+        if isinstance(first, tuple):
+            self.path = first + rest
+        elif "." in first:
+            self.path = tuple(first.split(".")) + rest
+        else:
+            self.path = (first,) + rest
+
+    def __iter__(self) -> Iterator[str]:
+        yield from self.path
+
+    def __hash__(self) -> int:
+        return hash(self.path)
+
+    def __repr__(self) -> str:
+        return ".".join(self)
+
+    def __add__(self, path: str) -> ConfigKey:
+        return ConfigKey(self.path, path)
+
+    def __len__(self) -> int:
+        return len(self.path)
+
+    def exists_in(self, source: Mapping | None) -> bool:
+        if source is None:
+            return False
+        doc = source
+        for key in self:
+            if key in doc:
+                doc = doc[key]
+            else:
+                return False
+        return True
+
+    def value_of(self, source: Mapping) -> Any:
+        doc = source
+        for key in self:
+            doc = doc[key]
+        return doc
+
+
+def iter_keys(
+    source: Mapping,
+    path: ConfigKey | None = None,
+) -> Iterator[ConfigKey]:
+    for key in source.keys():
+        if path is None:
+            local_path = ConfigKey(key)
+        else:
+            local_path = path + key
+        if isinstance(source[key], Mapping):
+            yield from iter_keys(source[key], local_path)
+        else:
+            yield local_path
+
+
+@define
+class Config:
+    paths: Paths = Factory(Paths)
+    credentials: Credentials = Factory(Credentials)
+    cli: Cli = Factory(Cli)
+    db: Db = Factory(Db)
+    download: Download = Factory(Download)
+    api: API = Factory(API)
+    _raw: TOMLDocument | None = field(init=False, default=None)
+    _config_file: Path | None = field(init=False, default=None)
+
+    @classmethod
+    def load(cls, path: Path) -> Config:
+        config_file = path / CONFIG_FILENAME
+        if config_file.is_file():
+            with config_file.open() as fh:
+                doc = tomlkit.load(fh)
+                for fixer in config_fixers:
+                    try:
+                        doc = fixer(doc)
+                    except Exception:
+                        raise BadConfigError(config_file)
+                raw = doc
+        else:
+            doc = TOMLDocument()
+            raw = None
+        config = _converter_defaults.structure(doc, cls)
+        config._raw = raw
+        config._config_file = config_file
+        return config
+
+    @classmethod
+    def default(cls) -> Config:
+        if InstallConfig.current is not None:
+            root = InstallConfig.current.path
+        else:
+            root = InstallConfig.default
+        return cls.load(root)
+
+    @property
+    def kind(self) -> ConfigKind:
+        if self._config_file is None:
+            return ConfigKind.Virtual
+        elif not self._config_file.is_file():
+            return ConfigKind.NoFile
+        elif self.unset_options():
+            return ConfigKind.Partial
+        else:
+            return ConfigKind.Complete
+
+    def dumps(self, defaults: bool = True, comments: bool = False) -> str:
+        return self.dump(defaults, comments).as_string()
+
+    def dump(
+        self,
+        defaults: bool = True,
+        comments: bool = False,
+    ) -> TOMLDocument:
+        if defaults:
+            converter = _converter_defaults
+        else:
+            converter = _converter_no_defaults
+        dump = converter.unstructure(self)
+        if not defaults:
+            pop_empty(dump)
+        doc = TOMLDocument()
+        doc.update(dump)
+        # if comments and self._raw is not None:
+        #     original = tomlkit.loads(self._raw)
+        return doc
+
+    def update_item(
+        self,
+        key: str,
+        value: int | str,
+        empty_ok: bool = False,
+    ) -> int | str | None:
+        if self._raw is None and empty_ok:
+            self._raw = TOMLDocument()
+        if self._raw is None:
+            raise VirtualConfigError
+        else:
+            doc: dict = self._raw
+        obj = self
+        *parts, last = ConfigKey(key)
+        for part in parts:
+            doc.setdefault(part, {})
+            doc = doc[part]
+            obj = getattr(obj, part)
+        old_value = getattr(obj, last)
+        if isinstance(old_value, str):
+            ...
+        elif isinstance(old_value, Container):
+            raise KeyError(key)
+        try:
+            value = int(value)
+        except ValueError:
+            ...
+        setattr(obj, last, value)
+        doc[last] = value
+        return old_value
+
+    def set_default(self, key: str) -> int | str | None:
+        ckey = ConfigKey(key)
+        if self._raw is None:
+            raise VirtualConfigError()
+        elif not ckey.exists_in(self._raw):
+            return None
+        default_config = self.__class__()
+        default_value = ckey.value_of(default_config.dump())
+        old_value: Any = ckey.value_of(self.dump())
+        first_pass = True
+        obj = self
+        for idx in range(len(ckey), 0, -1):
+            *parts, last = ckey.path[:idx]
+            doc: tomlkit.container.Container = self._raw
+            for part in parts:
+                if first_pass:
+                    obj = getattr(obj, part)
+                doc = cast(tomlkit.container.Container, doc[part])
+            if first_pass:
+                doc.remove(last)
+                setattr(obj, last, default_value)
+                first_pass = False
+            elif (
+                (value := doc[last])
+                and isinstance(value, tomlkit.container.Container)
+                and len(value) == 0
+            ):
+                doc.remove(last)
+        return old_value
+
+    def unset_options(self) -> list[ConfigKey]:
+        result: list[ConfigKey] = []
+        raw: dict
+        dump = self.dump()
+        if self._raw is None:
+            raw = {}
+        else:
+            raw = self._raw
+        for ckey in iter_keys(dump):
+            if not ckey.exists_in(raw):
+                result.append(ckey)
+        return result
+
+    def generate(
+        self,
+        overwrite: bool = False,
+    ) -> None:
+        match (self.kind, overwrite):
+            case (ConfigKind.Virtual, _):
+                raise VirtualConfigError
+            case (ConfigKind.Partial, overwrite):
+                defaults = self.dump()
+                for ckey in self.unset_options():
+                    self.update_item(str(ckey), ckey.value_of(defaults))
+            case (ConfigKind.Partial | ConfigKind.Complete, _):
+                raise FileExistsError(self._config_file)
+            case (ConfigKind.NoFile, _):
+                self._raw = self.dump()
+            case _:
+                raise ValueError(self.kind)
+        self.write()
+
+    def write(self) -> None:
+        if self.kind == ConfigKind.Virtual or self._raw is None:
+            raise VirtualConfigError
+        config_file = cast(Path, self._config_file)
+        with config_file.open("w") as f:
+            tomlkit.dump(self._raw, f)
+
+
+def _make_converter(omit_default: bool) -> Converter:
+    conv = Converter(omit_if_default=omit_default, forbid_extra_keys=True)
+    conv.register_unstructure_hook(Path, str)
+    conv.register_unstructure_hook(
+        Config,
+        make_dict_unstructure_fn(
+            Config,
+            conv,
+            _raw=override(omit=True),
+            _config_file=override(omit=True),
+        ),
+    )
+    return conv
+
+
+_converter_defaults = _make_converter(omit_default=False)
+_converter_no_defaults = _make_converter(omit_default=True)
+
+
+def pop_empty(d: dict[str, Any]) -> None:
+    keys = list(d.keys())
+    for key in keys:
+        value = d[key]
+        if isinstance(value, dict):
+            pop_empty(value)
+            if not value:
+                d.pop(key)

esgpull/constants.py

@@ -0,0 +1,56 @@
+CONFIG_FILENAME = "config.toml"
+ROOT_ENV = "ESGPULL_CURRENT"
+
+IDP = "/esgf-idp/openid/"
+CEDA_IDP = "/OpenID/Provider/server/"
+PROVIDERS = {
+    "esg-dn1.nsc.liu.se": IDP,
+    "esgf-data.dkrz.de": IDP,
+    "ceda.ac.uk": CEDA_IDP,
+    "esgf-node.ipsl.upmc.fr": IDP,
+    "esgf-node.llnl.gov": IDP,
+    "esgf.nci.org.au": IDP,
+}
+
+
+DEFAULT_FACETS = [
+    "project",
+    "mip_era",
+    "experiment",
+    "experiment_id",
+    "institute",
+    "institution_id",
+    "model",
+    "table_id",
+    "activity_id",
+    "ensemble",
+    "variant_label",
+    "realm",
+    "frequency",
+    "time_frequency",
+    "variable",
+    "variable_id",
+    "dataset_id",
+    "source_id",
+    "domain",
+    "driving_model",
+    "rcm_name",
+    "member_id",
+    "cmor_table",
+]
+EXTRA_FACETS = [
+    "query",
+    "start",
+    "end",
+    # "fields",
+    "facets",
+    "url",
+    "data_node",
+    "index_node",
+    "master_id",
+    "instance_id",  # search does not work with instance_id
+    "title",
+    "variable_long_name",
+    "experiment_family",
+]
+DEFAULT_CONSTRAINTS_WITH_VALUE: dict[str, str] = {}