46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
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))
|