from BTrees.IIBTree import IIBTree from ZODB import DB from ZODB.Connection import Connection from ZODB.FileStorage import FileStorage from persistent import Persistent from persistent.mapping import PersistentMapping from transaction import commit from telegram import Message, Chat def get_conn(read_only=False) -> Connection: storage = FileStorage('users.fs', read_only=read_only) db = DB(storage) conn = db.open() if not hasattr(conn.root, 'subscribers'): conn.root.subscribers = PersistentMapping() # migration 1 if not hasattr(conn.root, 'counter'): conn.root.counter = 0 for user in conn.root.subscribers.values(): # type: Subscriber if not hasattr(user, 'messages_forward') or not isinstance(user.messages_forward, IIBTree): user.messages_forward = IIBTree() user.messages_reverse = IIBTree() # end migrations commit() return conn class Subscriber(Persistent): def __init__(self, user_id, name): self.uid = user_id self.name = name self.messages_forward = IIBTree() self.messages_reverse = IIBTree() def update_from_message(self, m: Message): self.name = Subscriber.get_name(m.chat) @classmethod def get_name(cls, chat: Chat): return f'{chat.first_name or ""} {chat.last_name or ""} {chat.title or ""}'.strip() @classmethod def from_chat(cls, chat: Chat): return cls(chat.id, cls.get_name(chat))