from typing import Dict from pydantic import BaseModel, PrivateAttr from million.model.message import Message _default_file_path = './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(_default_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 msg.gifs[0].uri