data_comments.py

from collections import defaultdict

from data import load_jsonl

all_comments = load_jsonl('comments.jsonl')

def filter_comments(
    comments: list[dict] | None = None,
    experiment_id: int | list[int] | None = None,
    reason: str | list[str] | None = None,
    language: str | list[str] | None = None,
) -> list[dict]:
    filtered_comments = comments
    if filtered_comments is None:
        filtered_comments = all_comments
    if experiment_id is not None:
        if isinstance(experiment_id, int):
            experiment_id = {experiment_id}
        filtered_comments = [comment for comment in filtered_comments if comment['experiment_id'] in experiment_id]
    if reason is not None:
        if isinstance(reason, str):
            reason = {reason}
        filtered_comments = [comment for comment in filtered_comments if comment['reason'] in reason]
    if language is not None:
        if isinstance(language, str):
            language = {language}
        filtered_comments = [comment for comment in filtered_comments if comment['language'] in language]
    return filtered_comments