Messenger_1Million_Stats/million/analyze/word_finder.py

from datetime import date
import re
import million.analyze.message_evaluation as msg_val
from typing import List
from million.model.message import Message


def filter_words(messages: List[Message], words: List[str]) -> List[Message]:
    """
    Return every message containg the given words, you can use regex syntax inside your words
    i.e. find_words(messages, ["dogs?","m(ous|ic)e"])
    will search for : dog, dogs, mouse, mice
    """
    r_words = [rf"\b{w}\b" for w in words]
    rgx = "(" + "|".join(r_words) + ")"
    return [m for m in messages if m.content and re.search(rgx, m.content, re.I)]


def filter_value(messages: List[Message], val: int) -> List[Message]:
    """
    Return every message whose value is evaluated to the given val
    """
    return [m for m in messages if msg_val.get(m) == val]


def filter_date(messages: List[Message], other: date) -> List[Message]:
    """
    Return every message posted on the given date
    """
    return [m for m in messages if m.date_time.date() == other]


def filter_neighbours(
    messages: List[Message], msg: Message, amplitude: int = 10
) -> List[Message]:
    """
    Return the messages posted just before and after the given one.
    amplitude will indicate how many messages to look for in each direction
    """
    idx = messages.index(msg)
    start_index = max(0, idx - amplitude)
    end_index = min(len(messages), idx + amplitude + 1)
    return messages[start_index:end_index]


def retain_counts(messages: List[Message]) -> List[Message]:
    """
    Retain only the messages for which are considered having a counted value
    """
    return [msg for msg in messages if msg_val.get(msg)]