Source code for analyzer.cli.cli

import logging
from enum import Enum


import click
from rich import print



logger = logging.getLogger("analyzer")





def jumpIn(**kwargs):
    import code
    import readline
    import rlcompleter

    vars = globals()
    vars.update(locals())
    vars.update(kwargs)
    readline.set_completer(rlcompleter.Completer(vars).complete)
    readline.parse_and_bind("tab: complete")
    code.InteractiveConsole(vars).interact()





class LogLevel(str, Enum):


    ERROR = "ERROR"



    WARNING = "WARNING"



    INFO = "INFO"



    DEBUG = "DEBUG"




@click.group()
@click.option(
    "--log-level",
    default=LogLevel.INFO,
    type=click.Choice(LogLevel, case_sensitive=False),
)


def cli(log_level):
    logging.getLogger("analyzer").setLevel(log_level.value)



@cli.command()
@click.argument("input", type=click.Path(exists=True, file_okay=True, dir_okay=False))
@click.argument("output", type=click.Path(file_okay=False, dir_okay=True))
@click.option("--executor", "-e", type=str, required=True)
@click.option("--max-sample-events", type=int, default=None)
@click.option("--filter-dataset", type=str, default=None)
@click.option("--filter-sample", type=str, default=None)


def run(
    input,
    output,
    executor,
    max_sample_events,
    filter_dataset,
    filter_sample,
):
    from analyzer.core.running import runFromPath

    runFromPath(
        input,
        output,
        executor,
        max_sample_events=max_sample_events,
        filter_dataset=filter_dataset,
        filter_sample=filter_sample,
    )



@cli.command()
@click.argument(
    "config-path", type=click.Path(exists=True, file_okay=True, dir_okay=False)
)


def lint(config_path):
    from analyzer.core.analysis import loadAnalysis
    from analyzer.core.linting import runLint, LintLevel
    from rich import print
    import sys

    analysis = loadAnalysis(config_path)
    lint_msgs = runLint(analysis)

    if not lint_msgs:
        print("[bold green]Linting passed! No errors or warnings found.[/bold green]")
        return

    print("[bold yellow]Linter Warnings/Errors found:[/bold yellow]")
    for msg in lint_msgs:
        color = "red" if msg.level == LintLevel.ERROR else "yellow"
        print(f"[{color}]{msg}[/{color}]")

    if any(msg.level == LintLevel.ERROR for msg in lint_msgs):
        print(
            "[bold red]Analysis configuration failed linting checks with ERRORs.[/bold red]"
        )
        sys.exit(1)



@cli.command()
@click.argument("input", type=click.Path(exists=True, file_okay=True, dir_okay=False))


def glancelz4(input):
    import lz4.frame
    from rich import print
    import pickle as pkl

    with lz4.frame.open(input, "rb") as f:
        data = pkl.loads(f.read())
    print(data)



@cli.command()
@click.argument("config-path", type=str)
@click.argument("dataset-name", type=str)
@click.argument("sample-name", type=str)
@click.argument("event-source", type=str)
@click.argument("event-start", type=int)
@click.argument("event-stop", type=int)


def run_chunk(
    config_path, dataset_name, sample_name, event_source, event_start, event_stop
):
    from analyzer.core.running import getRepos
    from analyzer.core.event_collection import FileChunk
    from analyzer.core.analysis import loadAnalysis
    from analyzer.utils.structure_tools import getWithMeta

    analysis = loadAnalysis(config_path)
    dataset_repo, era_repo = getRepos(
        analysis.extra_dataset_paths, analysis.extra_era_paths
    )

    dataset = dataset_repo[dataset_name]
    sample, meta = getWithMeta(dataset, sample_name)
    meta = dict(meta)
    meta["era"] = era_repo[meta["era"]]
    chunk = FileChunk(
        file_path=config_path, event_start=event_start, event_stop=event_stop
    )
    analysis.analyzer.run(chunk, meta)



@cli.command()
@click.argument("config-path", type=str)
@click.argument("dataset-name", type=str)
@click.argument("sample-name", type=str)


def describe_analysis(config_path, dataset_name, sample_name):
    from analyzer.core.running import getRepos
    from analyzer.core.analysis import loadAnalysis
    from analyzer.utils.structure_tools import getWithMeta

    analysis = loadAnalysis(config_path)
    dataset_repo, era_repo = getRepos(
        analysis.extra_dataset_paths, analysis.extra_era_paths
    )

    dataset = dataset_repo[dataset_name]
    sample, meta = getWithMeta(dataset, sample_name)
    meta = dict(meta)
    print(analysis.analyzer)


    # output = analysis.analyzer.run(chunk, meta)


@cli.command()
@click.argument("files", nargs=-1)
@click.option("--configuration", "-c", type=str, required=str)
@click.option("--only-bad", is_flag=True)


def check(files, configuration, only_bad):
    from analyzer.cli.result_status import renderStatuses
    from analyzer.core.analysis import getSamples, loadAnalysis
    from analyzer.core.results import checkResults
    from analyzer.core.running import getRepos

    analysis = loadAnalysis(configuration)
    dataset_repo, era_repo = getRepos(
        analysis.extra_dataset_paths, analysis.extra_era_paths
    )
    all_samples = getSamples(analysis, dataset_repo)

    ret = checkResults(files)
    renderStatuses(ret, all_samples, only_bad=only_bad)



@cli.command()
@click.argument("inputs", type=str, nargs=-1)
@click.option("--output", "-o", type=str, required=True)
@click.option("--configuration", "-c", type=str, required=True)
@click.option("--executor", "-e", type=str, required=True)
@click.option("--filter-dataset", type=str, default=None)
@click.option("--filter-sample", type=str, default=None)


def patch(
    inputs,
    output,
    configuration,
    executor,
    filter_dataset,
    filter_sample,
):
    from analyzer.core.running import patchFromPath

    patchFromPath(
        configuration,
        inputs,
        output,
        executor,
        filter_dataset=filter_dataset,
        filter_sample=filter_sample,
    )



@cli.command()
@click.argument("inputs", type=str, nargs=-1)
@click.option("--interpretter", is_flag=True)
@click.option("--peek", is_flag=True, default=False)
@click.option("--merge-datasets", is_flag=True)


def browse(inputs, interpretter, peek, merge_datasets):
    from analyzer.core.results import loadResults, mergeAndScale
    from analyzer.core.serialization import converter, setupConverter

    setupConverter(converter)
    res = loadResults(inputs, peek_only=peek)
    if merge_datasets:
        res = mergeAndScale(res)
    if interpretter:
        jumpIn(results=res)
    else:
        from analyzer.cli.browser import ResultBrowser

        browser = ResultBrowser(res)
        browser.run()



@cli.command()
@click.argument("configuration", type=str, nargs=1)
@click.argument("inputs", type=str, nargs=-1)
@click.option("--prefix", type=str, required=False, default=None)
@click.option("--parallel", type=int, required=False, default=None)
@click.option("--target-load-size", type=int, required=False, default=None)
@click.option("--include-sidecar", is_flag=True, default=False)
@click.option(
    "--explain-grouping-only",
    is_flag=True,
    default=False,
    help="Dry-run the structure blocks and print a trace of how items are selected, grouped, and transformed.",
)
@click.option(
    "--explain-grouping-verbose",
    is_flag=True,
    default=False,
    help="",
)


def postprocess(
    configuration,
    inputs,
    parallel,
    prefix,
    target_load_size,
    include_sidecar,
    explain_grouping_only,
    explain_grouping_verbose,
):
    if explain_grouping_only:
        from analyzer.postprocessing.running import loadPostprocessor
        from analyzer.core.results import loadResults, mergeAndScale
        from analyzer.postprocessing.explain import renderTrace
        from rich import print as rprint

        postprocessor = loadPostprocessor(configuration)
        results = loadResults(inputs)
        if postprocessor.do_merge_and_scale:
            results = mergeAndScale(
                results, drop_sample_pattern=postprocessor.drop_sample_pattern
            )

        for proc_idx, processor in enumerate(postprocessor.processors):
            proc_name = type(processor).__name__
            rprint(f"\n[bold]━━━ Processor {proc_idx}: {proc_name} ━━━[/bold]")
            traces = processor.explain(results)
            for input_desc, trace in traces:
                tree = renderTrace(
                    trace,
                    label=f"inputs: {input_desc}",
                    verbose=explain_grouping_verbose,
                )
                rprint(tree)
        return

    from analyzer.postprocessing.running import runPostprocessors

    runPostprocessors(
        configuration,
        inputs,
        parallel=parallel,
        prefix=prefix,
        target_load_size=target_load_size,
        include_sidecar=include_sidecar,
    )



@cli.command()
@click.argument("inputs", type=str, nargs=-1)
@click.option(
    "--output", "-o", type=click.Path(file_okay=False, dir_okay=True), required=True
)
@click.option(
    "--group-by",
    type=click.Choice(["dataset", "sample", "none"], case_sensitive=False),
    default="dataset",
)


def merge(inputs, output, group_by):
    from analyzer.core.results import loadResults, ResultGroup
    from pathlib import Path

    print(f"Loading {len(inputs)} files...")
    res = loadResults(inputs)
    print("Merging done, saving...")

    if group_by == "dataset":
        output = Path(output)
        output.mkdir(exist_ok=True, parents=True)
        for dataset_name in res.keys():
            print(f"Saving {dataset_name}...")
            # Create a new ResultGroup with just this dataset
            ds_res = ResultGroup(
                name=res.name,
                results={dataset_name: res[dataset_name]},
                metadata=res.metadata,
            )
            with open(output / f"{dataset_name}.result", "wb") as f:
                f.write(ds_res.toBytes())
    elif group_by == "sample":
        output = Path(output)
        output.mkdir(exist_ok=True, parents=True)
        for dataset_name in res.keys():
            dataset = res[dataset_name]
            for sample_name in dataset.keys():
                print(f"Saving {dataset_name}__{sample_name}...")
                # Create nested ResultGroup structure
                sample_res = ResultGroup(
                    name=res.name,
                    results={
                        dataset_name: ResultGroup(
                            name=dataset.name,
                            results={sample_name: dataset[sample_name]},
                            metadata=dataset.metadata,
                        )
                    },
                    metadata=res.metadata,
                )
                with open(output / f"{dataset_name}__{sample_name}.result", "wb") as f:
                    f.write(sample_res.toBytes())
    elif group_by == "none":
        with open(output, "wb") as f:
            f.write(res.toBytes())



@cli.group()


def cache():
    pass



@cache.command()
@click.option("--tag", type=str, default=None)


def clear(tag):
    from analyzer.core.caching import cache

    if not click.confirm("Are you sure you want to clear the cache?"):
        return

    if tag is None:
        cache.clear()
    else:
        cache.evict(tag)



@cache.command()


def list():
    from analyzer.core.caching import cache

    for f in cache:
        print(f)



@cli.group("list")


def listData():
    pass



@click.option("--filter", type=str)
@click.option("--csv", is_flag=True)
@listData.command()


def samples(filter, csv):
    from analyzer.cli.dataset_table import createSampleTable
    from analyzer.core.running import getRepos
    from analyzer.utils.querying import Pattern
    from analyzer.core.serialization import converter, setupConverter

    setupConverter(converter)

    if filter:
        filter_pattern = Pattern(filter)
    else:
        filter_pattern = None
    dataset_repo, era_repo = getRepos()
    table = createSampleTable(dataset_repo, pattern=filter_pattern, as_csv=csv)
    print(table)



@click.option("--filter", type=str)
@click.option("--csv", is_flag=True)
@listData.command()


def datasets(filter, csv):
    from analyzer.cli.dataset_table import createDatasetTable
    from analyzer.core.running import getRepos
    from analyzer.utils.querying import Pattern
    from analyzer.core.serialization import converter, setupConverter

    setupConverter(converter)

    if filter:
        filter_pattern = Pattern(filter)
    else:
        filter_pattern = None
    dataset_repo, era_repo = getRepos()
    table = createDatasetTable(dataset_repo, pattern=filter_pattern, as_csv=csv)
    print(table)



@listData.group()


def eras():
    raise NotImplementedError()



@cli.command("search-modules")
@click.argument("query", type=str)


def searchModules(query):
    from analyzer.cli.module_search import searchModules as runSearch

    runSearch(query)



@cli.command("export-adl")
@click.argument("config-path", type=str)
@click.option(
    "--metadata",
    "-m",
    multiple=True,
    type=str,
    help="Metadata sets in format key=value,key=value. Example: era=2018,type=MC,name=2018_MC",
)
@click.option(
    "--output-dir", "-o", type=str, default=".", help="Directory to save the ADL files"
)
@click.option(
    "--ignore-module",
    multiple=True,
    type=str,
    help="Regex pattern of module names to ignore",
)


def exportAdl(config_path, metadata, output_dir, ignore_module):
    from analyzer.core.analysis import loadAnalysis
    from analyzer.core.adl import MetadataSpec, buildMetadata
    from analyzer.core.running import getRepos
    from pathlib import Path
    from analyzer.core.datasets import SampleType

    analysis = loadAnalysis(config_path)
    dataset_repo, era_repo = getRepos(
        analysis.extra_dataset_paths, analysis.extra_era_paths
    )

    outdir = Path(output_dir)
    outdir.mkdir(parents=True, exist_ok=True)

    ignore_pattern = None
    if ignore_module:
        ignore_pattern = "|".join(f"({m})" for m in ignore_module)

    if not metadata:
        print("Warning: No metadata specified. Using default 2018 MC context.")
        metadata = ["era=2018,type=MC,name=2018_MC"]

    for m_str in metadata:
        m_dict = dict(kv.split("=") for kv in m_str.split(","))
        era = m_dict.get("era", "2018")
        sample_type = m_dict.get("type", "MC")
        label = m_dict.get("name", f"{era}_{sample_type}")

        st = SampleType.MC if sample_type.upper() == "MC" else SampleType.Data
        spec = MetadataSpec(era=era, sampleType=st, label=label)

        full_meta = buildMetadata(spec, era_repo)

        adl_content = analysis.analyzer.exportAdl(
            metadata=full_meta,
            ignore_pattern=ignore_pattern,
            title=f"Analysis Export for {label}",
            config_path=config_path,
        )

        out_file = outdir / f"{label}.adl"
        with open(out_file, "w") as f:
            f.write(adl_content)
        print(f"Exported ADL to {out_file}")





def main():
    from analyzer.logging import setupLogging

    setupLogging()
    cli()