Messenger_1Million_Stats/million/analyze/dns_solver.py

from os.path import basename
from typing import Dict

from pydantic import BaseModel, PrivateAttr
from million.model.message import Message

_default_file_path = 'data/DNS'

class DNS_solver(BaseModel):
    file_path:str = _default_file_path

    _bank: Dict[str, int] | None = PrivateAttr(None)

    def solve(self, msg: Message) -> int:
        if self._bank == None: 
            self._bank = self.load(self.file_path)

        k = self._get_key(msg)
        if k and k in self._bank: return self._bank[k]

        return None

    def load(self, file_name: str) -> Dict[str, int]:
        result = {}
        with open(file_name, 'r') as f:
            for line in f:
                a,b = line.split()
                result[a] = int(b)

        return result

    def _get_key(self, msg: Message) -> str:

        # look into msg attributes
        # find uri
        return (msg.share or None) and msg.share.link or \
            (msg.gifs or None) and basename(msg.gifs[0].uri) or \
            (msg.photos or None) and basename(msg.photos[0].uri)