[ENH]: Add junifer list-elements for listing out elements of an YAML's DataGrabber #323

Merged
synchon merged 5 commits from feat/list-elements into main 2024-04-16 14:29:04 +00:00
5 changed files with 217 additions and 3 deletions

View file

@ -0,0 +1 @@
Add ``junifer list-elements`` to list out available elements for a DataGrabber based on filtering via ``--element`` by `Synchon Mandal`_

View file

@ -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."""

View file

@ -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)

View file

@ -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

View file

@ -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