annif.cli.run_index()

Complexity

Conditions

19

Size

Total Lines	134
Code Lines	98

Duplication

Lines	0
Ratio	0 %

Importance

Changes

0

"""Definitions for command-line (Click) commands for invoking Annif
@cli.command("index")
@cli_util.project_id
@click.argument("paths", type=click.Path(exists=True), nargs=-1)
@click.option(
    "--suffix", "-s", default=".annif", help="File name suffix for result files"
)
@click.option(
    "--force/--no-force",
    "-f/-F",
    default=False,
    help="Force overwriting of existing result files",
)
@click.option("--limit", "-l", default=10, help="Maximum number of subjects")
@click.option("--threshold", "-t", default=0.0, help="Minimum score threshold")
@click.option("--language", "-L", help="Language of subject labels")
@click.option(
    "--gzip/--no-gzip",
    "-z/-Z",
    "use_gzip",
    default=False,
    help="Gzip compress result files",
)
@click.option(
    "--output",
    "-O",
    type=click.Path(dir_okay=False, writable=True),
    default=None,
    help="Redirect all output to the given file (or '-' for stdout)",
)
@click.option(
    "--include-doc/--no-include-doc",
    "-i/-I",
    default=True,
    help="Include input documents in output",
)
@cli_util.backend_param_option
@cli_util.common_options
def run_index(
    project_id,
    paths,
    suffix,
    force,
    limit,
    threshold,
    language,
    backend_param,
    use_gzip=False,
    output=None,
    include_doc=True,
):
    """
    Index documents from directories or files, suggesting subjects for each document.
    Write the results in TSV files (for directories) or JSONL files (for files) with
    the given suffix (.jsonl suffix will be added to JSONL files).
    """
    project = cli_util.get_project(project_id)
    lang = language or project.vocab_lang
    if lang not in project.vocab.languages:
        raise click.BadParameter(f'language "{lang}" not supported by vocabulary')
    backend_params = cli_util.parse_backend_params(backend_param, project)

    # Helper function to process a directory
    def process_directory(directory):
        corpus = DocumentDirectory(directory, require_subjects=False)
        results = project.suggest_corpus(corpus, backend_params).filter(
            limit, threshold
        )

        for doc, suggestions in zip(corpus.documents, results):
            subjectfilename = re.sub(r"\.(txt|json)$", suffix, doc.file_path)
            if os.path.exists(subjectfilename) and not force:
                click.echo(
                    "Not overwriting {} (use --force to override)".format(
                        subjectfilename
                    )
                )
                continue
            with open(subjectfilename, "w", encoding="utf-8") as subjfile:
                cli_util.show_hits(suggestions, project, lang, file=subjfile)

    # Helper function to process a file
    def process_file(path):
        corpus = cli_util.open_doc_path(
            path, project.subjects, lang, require_subjects=False
        )
        results = project.suggest_corpus(corpus, backend_params).filter(
            limit, threshold
        )

        # Determine output stream
        if output == "-":
            stream_cm = cli_util.get_output_stream(path, suffix, "-", use_gzip, force)
        else:
            # For individual files, generate output filename based on input path
            if output:
                outfilename = output + (
                    ".gz" if use_gzip and not output.endswith(".gz") else ""
                )
                stream_cm = cli_util.get_output_stream(
                    path, suffix, outfilename, use_gzip, force
                )
            else:
                outfilename = (
                    re.sub(r"(\.[^.]+)?(\.gz)?$", "", path) + suffix + ".jsonl"
                )
                if use_gzip and not outfilename.endswith(".gz"):
                    outfilename += ".gz"
                stream_cm = cli_util.get_output_stream(
                    path, suffix, outfilename, use_gzip, force
                )

        if stream_cm is None:
            return

        with stream_cm as stream:
            for doc, suggestions in zip(corpus.documents, results):
                if include_doc:
                    output_data = doc.as_dict(project.subjects, lang)
                else:
                    output_data = {}
                    if doc.document_id:
                        output_data["document_id"] = doc.document_id
                output_data["results"] = [
                    suggestion_to_dict(suggestion, project.subjects, lang)
                    for suggestion in suggestions
                ]
                stream.write(json.dumps(output_data) + "\n")

    # Process paths in the order they were given
    for path in paths:
        if os.path.isdir(path):
            process_directory(path)
        elif os.path.isfile(path):
            process_file(path)