[ENH]: Add junifer list-elements for listing out elements of an YAML's DataGrabber #323
5 changed files with 217 additions and 3 deletions
1
docs/changes/newsfragments/323.feature
Normal file
1
docs/changes/newsfragments/323.feature
Normal file
|
|
@ -0,0 +1 @@
|
|||
Add ``junifer list-elements`` to list out available elements for a DataGrabber based on filtering via ``--element`` by `Synchon Mandal`_
|
||||
|
|
@ -8,7 +8,7 @@ import pathlib
|
|||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Tuple, Union
|
||||
from typing import Dict, List, Optional, Tuple, Union
|
||||
|
||||
import click
|
||||
import pandas as pd
|
||||
|
|
@ -20,6 +20,7 @@ from ..utils.logging import (
|
|||
warn_with_log,
|
||||
)
|
||||
from .functions import collect as api_collect
|
||||
from .functions import list_elements as api_list_elements
|
||||
from .functions import queue as api_queue
|
||||
from .functions import reset as api_reset
|
||||
from .functions import run as api_run
|
||||
|
|
@ -451,6 +452,69 @@ def reset(
|
|||
api_reset(config)
|
||||
|
||||
|
||||
@cli.command()
|
||||
@click.argument(
|
||||
"filepath",
|
||||
type=click.Path(
|
||||
exists=True, readable=True, dir_okay=False, path_type=pathlib.Path
|
||||
),
|
||||
)
|
||||
@click.option("--element", type=str, multiple=True)
|
||||
@click.option(
|
||||
"-o",
|
||||
"--output-file",
|
||||
type=click.Path(dir_okay=False, writable=True, path_type=pathlib.Path),
|
||||
)
|
||||
@click.option(
|
||||
"-v",
|
||||
"--verbose",
|
||||
type=click.UNPROCESSED,
|
||||
callback=_validate_verbose,
|
||||
default="info",
|
||||
)
|
||||
def list_elements(
|
||||
filepath: click.Path,
|
||||
element: Tuple[str],
|
||||
output_file: Optional[click.Path],
|
||||
verbose: Union[str, int],
|
||||
) -> None:
|
||||
"""Element listing command for CLI.
|
||||
|
||||
\f
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filepath : click.Path
|
||||
The filepath to the configuration file.
|
||||
element : tuple of str
|
||||
The element to operate on.
|
||||
output_file : click.Path or None
|
||||
The path to write the output to. If not None, writing to
|
||||
stdout is not performed.
|
||||
verbose : click.Choice
|
||||
The verbosity level: warning, info or debug (default "info").
|
||||
|
||||
"""
|
||||
configure_logging(level=verbose)
|
||||
# Parse YAML
|
||||
config = parse_yaml(filepath) # type: ignore
|
||||
# Fetch datagrabber
|
||||
datagrabber = config["datagrabber"]
|
||||
# Parse elements
|
||||
elements = _parse_elements(element, config)
|
||||
# Perform operation
|
||||
listed_elements = api_list_elements(
|
||||
datagrabber=datagrabber,
|
||||
elements=elements,
|
||||
)
|
||||
# Check if output file is provided
|
||||
if output_file is not None:
|
||||
output_file.touch()
|
||||
output_file.write_text(listed_elements)
|
||||
else:
|
||||
click.secho(listed_elements, fg="blue")
|
||||
|
||||
|
||||
@cli.group()
|
||||
def setup() -> None: # pragma: no cover
|
||||
"""Configure commands for Junifer."""
|
||||
|
|
|
|||
|
|
@ -361,3 +361,44 @@ def reset(config: Dict) -> None:
|
|||
shutil.rmtree(job_dir)
|
||||
# Remove directory
|
||||
job_dir.parent.rmdir()
|
||||
|
||||
|
||||
def list_elements(
|
||||
datagrabber: Dict,
|
||||
elements: Union[str, List[Union[str, Tuple]], Tuple, None] = None,
|
||||
) -> str:
|
||||
"""List elements of the datagrabber filtered using `elements`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
datagrabber : dict
|
||||
DataGrabber to index. Must have a key ``kind`` with the kind of
|
||||
DataGrabber to use. All other keys are passed to the DataGrabber
|
||||
constructor.
|
||||
elements : str or tuple or list of str or tuple, optional
|
||||
Element(s) to filter using. Will be used to index the DataGrabber
|
||||
(default None).
|
||||
|
||||
"""
|
||||
# Get datagrabber to use
|
||||
datagrabber_object = _get_datagrabber(datagrabber)
|
||||
|
||||
# Fetch elements
|
||||
raw_elements_to_list = []
|
||||
with datagrabber_object:
|
||||
if elements is not None:
|
||||
for element in datagrabber_object.filter(elements):
|
||||
raw_elements_to_list.append(element)
|
||||
else:
|
||||
for element in datagrabber_object:
|
||||
raw_elements_to_list.append(element)
|
||||
|
||||
elements_to_list = []
|
||||
for element in raw_elements_to_list:
|
||||
# Stringify elements if tuple for operation
|
||||
str_element = (
|
||||
",".join(element) if isinstance(element, tuple) else element
|
||||
)
|
||||
elements_to_list.append(str_element)
|
||||
|
||||
return "\n".join(elements_to_list)
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from ruamel.yaml import YAML
|
|||
from junifer.api.cli import (
|
||||
_parse_elements_file,
|
||||
collect,
|
||||
list_elements,
|
||||
queue,
|
||||
reset,
|
||||
run,
|
||||
|
|
@ -297,6 +298,88 @@ def test_reset(
|
|||
assert reset_result.exit_code == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"elements",
|
||||
[
|
||||
("sub-01", "sub-02"),
|
||||
("sub-03", "sub-04"),
|
||||
],
|
||||
)
|
||||
def test_list_elements_stdout(
|
||||
elements: Tuple[str, ...],
|
||||
) -> None:
|
||||
"""Test elements listing to stdout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
elements : tuple of str
|
||||
The parametrized elements for filtering.
|
||||
|
||||
"""
|
||||
# Get test config
|
||||
infile = Path(__file__).parent / "data" / "partly_cloudy_agg_mean_tian.yml"
|
||||
# List elements command arguments
|
||||
list_elements_args = [
|
||||
str(infile.absolute()),
|
||||
"--verbose",
|
||||
"debug",
|
||||
"--element",
|
||||
elements[0],
|
||||
"--element",
|
||||
elements[1],
|
||||
]
|
||||
# Invoke list elements command
|
||||
list_elements_result = runner.invoke(list_elements, list_elements_args)
|
||||
# Check
|
||||
assert list_elements_result.exit_code == 0
|
||||
assert f"{elements[0]}\n{elements[1]}" in list_elements_result.stdout
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"elements",
|
||||
[
|
||||
("sub-01", "sub-02"),
|
||||
("sub-03", "sub-04"),
|
||||
],
|
||||
)
|
||||
def test_list_elements_output_file(
|
||||
tmp_path: Path,
|
||||
elements: Tuple[str, ...],
|
||||
) -> None:
|
||||
"""Test elements listing to output file.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
elements : tuple of str
|
||||
The parametrized elements for filtering.
|
||||
|
||||
"""
|
||||
# Get test config
|
||||
infile = Path(__file__).parent / "data" / "partly_cloudy_agg_mean_tian.yml"
|
||||
# Output file
|
||||
output_file = tmp_path / "elements.txt"
|
||||
# List elements command arguments
|
||||
list_elements_args = [
|
||||
str(infile.absolute()),
|
||||
"--verbose",
|
||||
"debug",
|
||||
"--element",
|
||||
elements[0],
|
||||
"--element",
|
||||
elements[1],
|
||||
"--output-file",
|
||||
str(output_file.resolve()),
|
||||
]
|
||||
# Invoke list elements command
|
||||
list_elements_result = runner.invoke(list_elements, list_elements_args)
|
||||
# Check
|
||||
assert list_elements_result.exit_code == 0
|
||||
with open(output_file) as f:
|
||||
assert f"{elements[0]}\n{elements[1]}" == f.read()
|
||||
|
||||
|
||||
def test_wtf_short() -> None:
|
||||
"""Test short version of wtf command."""
|
||||
# Invoke wtf command
|
||||
|
|
|
|||
|
|
@ -7,13 +7,13 @@
|
|||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Tuple, Union
|
||||
from typing import Dict, List, Optional, Tuple, Union
|
||||
|
||||
import pytest
|
||||
from ruamel.yaml import YAML
|
||||
|
||||
import junifer.testing.registry # noqa: F401
|
||||
from junifer.api.functions import collect, queue, reset, run
|
||||
from junifer.api.functions import collect, list_elements, queue, reset, run
|
||||
from junifer.datagrabber.base import BaseDataGrabber
|
||||
from junifer.pipeline.registry import build
|
||||
|
||||
|
|
@ -637,3 +637,28 @@ def test_reset_queue(
|
|||
|
||||
assert not Path(storage["uri"]).exists()
|
||||
assert not (tmp_path / "junifer_jobs" / job_name).exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"elements",
|
||||
[
|
||||
["sub-01"],
|
||||
None,
|
||||
],
|
||||
)
|
||||
def test_list_elements(
|
||||
datagrabber: Dict[str, str],
|
||||
elements: Optional[List[str]],
|
||||
) -> None:
|
||||
"""Test elements listing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
datagrabber : dict
|
||||
Testing datagrabber as dictionary.
|
||||
elements : str of list of str
|
||||
The parametrized elements for filtering.
|
||||
|
||||
"""
|
||||
listed_elements = list_elements(datagrabber, elements)
|
||||
assert "sub-01" in listed_elements
|
||||
|
|
|
|||
Loading…
Reference in a new issue