Messenger_1Million_Stats/million/model/sequence.py

from __future__ import annotations

from pydantic import BaseModel
import pydantic

from million.model.message import Message
import million.analyze.message_evaluation as msg_val


class Sequence(BaseModel):
    start_message: Message
    end_message: Message | None = None

    @pydantic.validator('end_message', pre=True, always=True)
    def default_end_message(cls, v, *, values):
        return v or values['start_message'] 

    def start(self) -> int:
        return msg_val.get(self.start_message)
    
    def end(self) -> int:
        return msg_val.get(self.end_message)
    
    def length(self) -> int:
        return self.end() - self.start() + 1
    
    def merge(self, other: Sequence) -> None:
        if other.start() < self.start(): self.start_message = other.start_message
        if other.end() > self.end(): self.end_message = other.end_message
    
    def overlaps(self, other: Sequence) -> bool:
        return self.start() <= other.end() + 1 and \
            other.start() <= self.end() + 1