annif.vocab.subject_index

Complexity

Total Complexity

39

Size/Duplication

Total Lines	170
Duplicated Lines	0 %

Importance

Changes

0

20 Methods

Rating	Name	Duplication	Size	Complexity
A	SubjectIndexFile.load_subjects()	0	6	2
A	SubjectIndexFile.languages()	0	3	1
A	SubjectIndexFile.__init__()	0	5	1
A	SubjectIndexFile.__len__()	0	2	1
A	SubjectIndexFilter.by_uri()	0	7	2
A	SubjectIndexFilter.contains_uri()	0	4	2
A	SubjectIndexFile.save()	0	14	5
A	SubjectIndexFilter.by_label()	0	9	3
A	SubjectIndexFilter.languages()	0	3	1
A	SubjectIndexFile.by_label()	0	6	2
A	SubjectIndexFilter.active()	0	9	1
A	SubjectIndexFile.__getitem__()	0	5	2
A	SubjectIndexFile.contains_uri()	0	2	1
A	SubjectIndexFile.load()	0	8	1
A	SubjectIndexFilter.__getitem__()	0	5	2
A	SubjectIndexFile.active()	0	6	1
A	SubjectIndexFilter.__init__()	0	3	1
A	SubjectIndexFile.append()	0	10	5
A	SubjectIndexFile.by_uri()	0	10	4
A	SubjectIndexFilter.__len__()	0	2	1

"""Subject index functionality for Annif"""

from __future__ import annotations

import csv
from typing import TYPE_CHECKING

import annif
import annif.corpus
import annif.util

from .types import SubjectIndex

if TYPE_CHECKING:

    from annif.corpus.subject import Subject, SubjectCorpus


logger = annif.logger
logger.addFilter(annif.util.DuplicateFilter())


class SubjectIndexFile(SubjectIndex):
    """SubjectIndex implementation backed by a file."""

    def __init__(self) -> None:
        self._subjects = []
        self._uri_idx = {}
        self._label_idx = {}
        self._languages = None

    def load_subjects(self, corpus: SubjectCorpus) -> None:
        """Initialize the subject index from a subject corpus"""

        self._languages = corpus.languages
        for subject in corpus.subjects:
            self.append(subject)

    def __len__(self) -> int:
        return len(self._subjects)

    @property
    def languages(self) -> list[str] | None:
        return self._languages

    def __getitem__(self, subject_id: int) -> Subject:
        subject = self._subjects[subject_id]
        if subject.labels is None:
            raise IndexError(f"Subject is deprecated: {subject_id}")
        return subject

    def append(self, subject: Subject) -> None:
        if self._languages is None and subject.labels is not None:
            self._languages = list(subject.labels.keys())

        subject_id = len(self._subjects)
        self._uri_idx[subject.uri] = subject_id
        if subject.labels:
            for lang, label in subject.labels.items():
                self._label_idx[(label, lang)] = subject_id
        self._subjects.append(subject)

    def contains_uri(self, uri: str) -> bool:
        return uri in self._uri_idx

    def by_uri(self, uri: str, warnings: bool = True) -> int | None:
        try:
            subject_id = self._uri_idx[uri]
            if self._subjects[subject_id].labels is None:  # deprecated
                return None
            return subject_id
        except KeyError:
            if warnings:
                logger.warning("Unknown subject URI <%s>", uri)
            return None

    def by_label(self, label: str | None, language: str) -> int | None:
        try:
            return self._label_idx[(label, language)]
        except KeyError:
            logger.warning('Unknown subject label "%s"@%s', label, language)
            return None

    @property
    def active(self) -> list[tuple[int, Subject]]:
        return [
            (subj_id, subject)
            for subj_id, subject in enumerate(self._subjects)
            if subject.labels is not None
        ]

    def save(self, path: str) -> None:
        """Save this subject index into a file with the given path name."""

        fieldnames = ["uri", "notation"] + [f"label_{lang}" for lang in self._languages]

        with open(path, "w", encoding="utf-8", newline="") as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            writer.writeheader()
            for subject in self:
                row = {"uri": subject.uri, "notation": subject.notation or ""}
                if subject.labels:
                    for lang, label in subject.labels.items():
                        row[f"label_{lang}"] = label
                writer.writerow(row)

    @classmethod
    def load(cls, path: str) -> SubjectIndex:
        """Load a subject index from a CSV file and return it."""

        vocab_file = annif.vocab.VocabFileCSV(path)
        subject_index = cls()
        subject_index.load_subjects(vocab_file)
        return subject_index


class SubjectIndexFilter(SubjectIndex):
    """SubjectIndex implementation that filters another SubjectIndex based
    on a list of subject URIs to exclude."""

    def __init__(self, subject_index: SubjectIndex, exclude: list[str]):
        self._subject_index = subject_index
        self._exclude = set(exclude)

    def __len__(self) -> int:
        return len(self._subject_index)

    @property
    def languages(self) -> list[str] | None:
        return self._subject_index.languages

    def __getitem__(self, subject_id: int) -> Subject:
        subject = self._subject_index[subject_id]
        if subject.uri in self._exclude:
            raise IndexError(f"Subject is excluded: {subject.uri}")
        return subject

    def contains_uri(self, uri: str) -> bool:
        if uri in self._exclude:
            return False
        return self._subject_index.contains_uri(uri)

    def by_uri(self, uri: str, warnings: bool = True) -> int | None:
        """return the subject ID of a subject by its URI, or None if not found.
        If warnings=True, log a warning message if the URI cannot be found."""

        if uri in self._exclude:
            return None
        return self._subject_index.by_uri(uri, warnings)

    def by_label(self, label: str | None, language: str) -> int | None:
        """return the subject ID of a subject by its label in a given
        language"""

        subject_id = self._subject_index.by_label(label, language)
        subject = self._subject_index[subject_id]
        if subject is not None and subject.uri not in self._exclude:
            return subject_id
        return None

    @property
    def active(self) -> list[tuple[int, Subject]]:
        """return a list of (subject_id, Subject) tuples of all subjects that
        are available for use"""

        return [
            (subject_id, subject)
            for subject_id, subject in self._subject_index.active
            if subject.uri not in self._exclude
        ]