annif.transform.select

Complexity

Total Complexity

5

Size/Duplication

Total Lines	34
Duplicated Lines	0 %

Importance

Changes

0

3 Methods

Rating	Name	Duplication	Size	Complexity
A	SelectTransform._get_texts()	0	6	3
A	SelectTransform.__init__()	0	3	1
A	SelectTransform.transform_doc()	0	4	1

"""A transformation that allows selecting Document metadata fields to be
used instead of the main text."""

from __future__ import annotations

from typing import TYPE_CHECKING

from annif.corpus import Document

from . import transform

if TYPE_CHECKING:
    from annif.project import AnnifProject


class SelectTransform(transform.BaseTransform):
    name = "select"

    def __init__(self, project: AnnifProject | None, *fields: str) -> None:
        super().__init__(project)
        self.fields = fields

    def _get_texts(self, doc):
        for fld in self.fields:
            if fld == "text":
                yield doc.text
            else:
                yield doc.metadata[fld]

    def transform_doc(self, doc: Document) -> Document:
        new_text = "\n".join(self._get_texts(doc))
        return Document(
            text=new_text, subject_set=doc.subject_set, metadata=doc.metadata
        )