You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

media_count_mapper.py 1.1KB

1234567891011121314151617181920212223242526272829303132333435363738
  1. from os.path import basename
  2. from typing import Dict
  3. from pydantic import BaseModel, PrivateAttr
  4. from million.model.message import Message
  5. _default_file_path = 'data/DefaultMediaCountMapFile'
  6. class MediaCountMapper(BaseModel):
  7. file_path:str = _default_file_path
  8. _bank: Dict[str, int] | None = PrivateAttr(None)
  9. def solve(self, msg: Message) -> int:
  10. if self._bank == None:
  11. self._bank = self.load(self.file_path)
  12. k = self._get_key(msg)
  13. if k and k in self._bank: return self._bank[k]
  14. return None
  15. def load(self, file_name: str) -> Dict[str, int]:
  16. result = {}
  17. with open(file_name, 'r') as f:
  18. for line in f:
  19. a,b = line.split()
  20. result[a] = int(b)
  21. return result
  22. def _get_key(self, msg: Message) -> str:
  23. # look into msg attributes
  24. # find uri
  25. return (msg.share or None) and msg.share.link or \
  26. (msg.gifs or None) and basename(msg.gifs[0].uri) or \
  27. (msg.photos or None) and basename(msg.photos[0].uri)