Переход на SQLite.

Добавлены команды для учета предупреждений.

.gitignore

@@ -1,3 +1,4 @@
 .idea
 .venv
-db.json
+bot.db
+config.json

bot.py

@@ -7,25 +7,26 @@ from vkbottle.bot import Bot, Message
 from vkbottle_types.objects import MessagesGetConversationMembers
 
 import config
-from config import db_load, db_save
+from config import config_load
+import database
 
-db_load()
-db_save()
-bot = Bot(config.DB['api_token'])
+
+config_load()
+bot = Bot(config.Config['api_token'])
 
 
 def posix_time():
     return calendar.timegm(time.gmtime())
 
 
-async def get_user_info(user_id):
+async def vk_get_user_info(user_id):
     info = await bot.api.users.get(user_ids=[user_id])
     if len(info) > 0:
         return info[0]
     return None
 
 
-def user_is_admin(user_id: int, chat_members: MessagesGetConversationMembers):
+def vk_user_is_admin(user_id: int, chat_members: MessagesGetConversationMembers):
     for member in chat_members.items:
         if member.member_id != user_id:
             continue
@@ -33,51 +34,80 @@ def user_is_admin(user_id: int, chat_members: MessagesGetConversationMembers):
     return False
 
 
+def create_user_if_not_exists(chat_id: str, user_id: str):
+    res = database.DB.get_user(chat_id, user_id)
+    if res is None:
+        database.DB.add_user(chat_id, user_id)
+
+
 @bot.on.chat_message(text="!старт")
-async def stats_handler(message: Message):
+async def start_handler(message: Message):
     chat_members = await bot.api.messages.get_conversation_members(peer_id=message.peer_id, extended=False)
 
-    if not user_is_admin(message.from_id, chat_members):
+    if not vk_user_is_admin(message.from_id, chat_members):
         await message.answer('О таком меня может попросить только администратор беседы.')
         return
 
-    chat_id = str(message.peer_id)
-    if chat_id not in config.DB['chats']:
-        config.DB['chats'][chat_id] = {'users': {}}
-
     for member in chat_members.items:
+        # Пропустить ботов
         if member.member_id < 0:
             continue
-        member_id = str(member.member_id)
-        if member_id not in config.DB['chats'][chat_id]['users']:
-            config.DB['chats'][chat_id]['users'][member_id] = {
-                'last_message': 0,
-                'messages_today': 0,
-                'messages_month': 0
-            }
-
-    db_save()
+        create_user_if_not_exists(message.peer_id, member.member_id)
 
     await message.answer('Готова к работе!')
 
 
+@bot.on.chat_message(text="!предупреждение")
+async def warning_handler(message: Message):
+    chat_members = await bot.api.messages.get_conversation_members(peer_id=message.peer_id, extended=False)
+
+    if not vk_user_is_admin(message.from_id, chat_members):
+        await message.answer('О таком меня может попросить только администратор беседы.')
+        return
+
+    if message.reply_message is None:
+        await message.answer('Эту команду нужно вызывать в ответном сообщении.')
+        return
+
+    chat_id = message.peer_id
+    user_id = message.reply_message.from_id
+    create_user_if_not_exists(chat_id, user_id)
+
+    database.DB.user_increment_warnings(chat_id, user_id)
+    user = database.DB.get_user(chat_id, user_id)
+
+    user_info = await vk_get_user_info(user_id)
+    if user_info is not None:
+        await message.answer('У участника {} {} {} предупреждений.'.format(
+            user_info.first_name,
+            user_info.last_name,
+            user['warnings'])
+        )
+
+
+@bot.on.chat_message(text="!предупреждения")
+async def warnings_handler(message: Message):
+    top_users = database.DB.get_top_warnings(message.peer_id)
+
+    response = '* Участники с предупреждениями:\n'
+    i = 1
+    for user in top_users:
+        info = await vk_get_user_info(user['user_id'])
+        response += '{}. {} {} - {}\n'.format(i, info.first_name, info.last_name, user['warnings'])
+        i = i + 1
+
+    await message.answer(response)
+
+
 @bot.on.chat_message(text="!сегодня")
 async def stats_today_handler(message: Message):
-    chat_id = str(message.peer_id)
+    top_users = database.DB.get_top_messages_today(message.peer_id)
 
-    chat_users = []
-    for user_id in config.DB['chats'][chat_id]['users'].keys():
-        flat_user = {**{'id': user_id}, **config.DB['chats'][chat_id]['users'][user_id]}
-        chat_users.append(flat_user)
-
-    top = sorted(chat_users, key=lambda item: item['messages_today'], reverse=True)
     response = '* Статистика за сегодня:\n'
-    i = 0
-    while i < len(top):
-        if top[i]['messages_today'] == 0:
-            break
-        info = await get_user_info(top[i]['id'])
-        response += '{}. {} {} - {}\n'.format(i + 1, info.first_name, info.last_name, top[i]['messages_today'])
+    i = 1
+    for user in top_users:
+        info = await vk_get_user_info(user['user_id'])
+        response += '{}. {} {} - {}\n'.format(i, info.first_name, info.last_name, user['messages_today'])
         i = i + 1
 
     await message.answer(response)
@@ -85,21 +115,13 @@ async def stats_today_handler(message: Message):
 
 @bot.on.chat_message(text="!месяц")
 async def stats_month_handler(message: Message):
-    chat_id = str(message.peer_id)
+    top_users = database.DB.get_top_messages_month(message.peer_id)
 
-    chat_users = []
-    for user_id in config.DB['chats'][chat_id]['users'].keys():
-        flat_user = {**{'id': user_id}, **config.DB['chats'][chat_id]['users'][user_id]}
-        chat_users.append(flat_user)
-
-    top = sorted(chat_users, key=lambda item: item['messages_month'], reverse=True)
-    response = '\n* Статистика за месяц:\n'
-    i = 0
-    while i < len(top):
-        if top[i]['messages_month'] == 0:
-            break
-        info = await get_user_info(top[i]['id'])
-        response += '{}. {} {} - {}\n'.format(i + 1, info.first_name, info.last_name, top[i]['messages_month'])
+    response = '* Статистика за месяц:\n'
+    i = 1
+    for user in top_users:
+        info = await vk_get_user_info(user['user_id'])
+        response += '{}. {} {} - {}\n'.format(i, info.first_name, info.last_name, user['messages_month'])
         i = i + 1
 
     await message.answer(response)
@@ -107,60 +129,34 @@ async def stats_month_handler(message: Message):
 
 @bot.on.chat_message(text="!молчуны")
 async def silent_handler(message: Message):
-    chat_id = str(message.peer_id)
-
-    chat_users = []
-    for user_id in config.DB['chats'][chat_id]['users'].keys():
-        flat_user = {**{'id': user_id},
-                     **{'last_message': config.DB['chats'][chat_id]['users'][user_id]['last_message']}}
-        chat_users.append(flat_user)
-    top = sorted(chat_users, key=lambda item: item['last_message'])
-
     now = posix_time()
+    threshold = now - 14 * 24 * 3600
+    top_users = database.DB.get_top_silent(message.peer_id, threshold)
 
     response = '* Молчуны чата (не писали больше 2 недель):\n'
-    for i in range(0, len(top)):
-        if top[i]['last_message'] == 0:
-            info = await get_user_info(top[i]['id'])
-            if info is None:
-                continue
-            response += '{}. {} {} - никогда\n'.format(i + 1, info.first_name, info.last_name)
+    i = 1
+    for user in top_users:
+        info = await vk_get_user_info(user['user_id'])
+        if user['last_message'] == 0:
+            response += '{}. {} {} - никогда\n'.format(i, info.first_name, info.last_name)
         else:
-            days_silent = round((now - top[i]['last_message']) / 3600 / 24)
-            if days_silent < 14:
-                break
-            info = await get_user_info(top[i]['id'])
-            response += '{}. {} {} - {} дней\n'.format(i + 1, info.first_name, info.last_name, days_silent)
+            days_silent = round((now - user['last_message']) / 3600 / 24)
+            response += '{}. {} {} - {} дней\n'.format(i, info.first_name, info.last_name, days_silent)
 
     await message.answer(response)
 
 
 @bot.on.chat_message()
 async def any_message_handler(message: Message):
+    # Игнорировать ботов
     if message.from_id < 0:
         return
 
-    chat_id = str(message.peer_id)
-    user_id = str(message.from_id)
-    if chat_id not in config.DB['chats']:
-        config.DB['chats'][chat_id] = {'users': {}}
-    if user_id not in config.DB['chats'][chat_id]['users']:
-        config.DB['chats'][chat_id]['users'][user_id] = {}
-
-    config.DB['chats'][chat_id]['users'][user_id]['last_message'] = posix_time()
-    config.DB['chats'][chat_id]['users'][user_id]['messages_today'] = (
-            config.DB['chats'][chat_id]['users'][user_id].get('messages_today', 0) + 1)
-    config.DB['chats'][chat_id]['users'][user_id]['messages_month'] = (
-            config.DB['chats'][chat_id]['users'][user_id].get('messages_month', 0) + 1)
-
-    db_save()
-
-
-# @bot.on.raw_event(GroupEventType.GROUP_JOIN, GroupTypes.GroupJoin)
-# async def group_join_handler(event: GroupTypes.GroupJoin):
-#     chat_id = str(event.group_id)
-#     user_id = str(event.object.user_id)
-#     if user_id == bot.api.users.get()
+    chat_id = message.peer_id
+    user_id = message.from_id
+    create_user_if_not_exists(chat_id, user_id)
+    database.DB.user_set_last_message(chat_id, user_id, posix_time())
+    database.DB.user_increment_messages(chat_id, user_id)
 
 
 async def wait_until(target_time: datetime.datetime):
@@ -187,17 +183,12 @@ async def counters_reset_task():
         await wait_until(target_datetime)
 
         print('Resetting daily counters...')
-        for chat_id in config.DB['chats']:
-            for user_id in config.DB['chats'][chat_id]['users']:
-                config.DB['chats'][chat_id]['users'][user_id]['messages_today'] = 0
+        database.DB.reset_messages_today()
 
         if target_datetime.day == 1:
             print('Resetting monthly counters...')
-            for chat_id in config.DB['chats']:
-                for user_id in config.DB['chats'][chat_id]['users']:
-                    config.DB['chats'][chat_id]['users'][user_id]['messages_month'] = 0
+            database.DB.reset_messages_month()
 
-        db_save()
         target_datetime = target_datetime + datetime.timedelta(days=1)
 
 

config.py

@@ -1,39 +1,23 @@
 import json
 
 
-# db = {
-#   chats: {
-#     CHAT_ID: {
-#       users: {
-#         USER_ID: {
-#           messages_today: 10
-#           messages_in_month: 1045
-#           last_message: 1755627442
-#         }
-#       }
-#     }
-#   },
+# Config = {
 #   api_token: "xxxxxxxx"
 # }
-DB = {}
+Config = {}
 
 
-def db_load():
-    global DB
+def config_load():
+    global Config
     try:
-        file = open('db.json', 'r')
-        DB = json.load(file)
+        file = open('config.json', 'r')
+        Config = json.load(file)
         file.close()
     except IOError:
-        DB = {'chats': {}, 'api_token': ''}
+        Config = {'api_token': ''}
 
 
-def db_save():
-    global DB
-    with open('db.json', 'w') as file:
-        json.dump(DB, file, indent=4)
-
-
-def db_print():
-    global DB
-    print(DB)
+def config_save():
+    global Config
+    with open('config.json', 'w') as file:
+        json.dump(Config, file, indent=4)

database.py

@@ -0,0 +1,99 @@
+import sqlite3
+from typing import List
+
+
+class Database:
+    def __init__(self):
+        self.conn = sqlite3.connect('bot.db')
+        self.conn.row_factory = sqlite3.Row
+        self.cursor = self.conn.cursor()
+
+        self.cursor.execute("""
+        CREATE TABLE IF NOT EXISTS users (
+            "chat_id"	INTEGER,
+            "user_id"	INTEGER,
+            "last_message"	INTEGER NOT NULL DEFAULT 0,
+            "messages_today"	INTEGER NOT NULL DEFAULT 0,
+            "messages_month"	INTEGER NOT NULL DEFAULT 0,
+            "warnings"	INTEGER NOT NULL DEFAULT 0,
+            PRIMARY KEY("chat_id","user_id"))
+        """)
+        self.conn.commit()
+
+    def get_user(self, chat_id: int, user_id: int):
+        self.cursor.execute("SELECT * FROM users WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
+        return self.cursor.fetchone()
+
+    def get_users(self, chat_id: int):
+        self.cursor.execute("SELECT * FROM users WHERE chat_id = ?", (chat_id,))
+        return self.cursor.fetchall()
+
+    def add_user(self, chat_id: int, user_id: int):
+        self.cursor.execute("INSERT INTO users (chat_id, user_id) VALUES (?, ?)", (chat_id, user_id))
+        self.conn.commit()
+
+    def user_set_last_message(self, chat_id: int, user_id: int, last_message: int):
+        self.user_update(chat_id, user_id, last_message=last_message)
+
+    def user_increment_messages(self, chat_id: int, user_id: int):
+        self.user_increment(chat_id, user_id, ['messages_today', 'messages_month'])
+        self.conn.commit()
+
+    def user_increment_warnings(self, chat_id: int, user_id: int):
+        self.user_increment(chat_id, user_id, ['warnings'])
+
+    def user_increment(self, chat_id: int, user_id: int, fields: List[str]):
+        self.cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) +
+                            " WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
+        self.conn.commit()
+
+    def user_update(self, chat_id: int, user_id: int, **kwargs):
+        # query = ("UPDATE users SET " + ", ".join(f + " = :" + f for f in kwargs) +
+        #         " WHERE chat_id = ? AND user_id = ?")
+        # print(query)
+        self.cursor.execute("UPDATE users SET " + ", ".join(f + " = :" + f for f in kwargs) +
+                            f" WHERE chat_id = {chat_id} AND user_id = {user_id}", kwargs)
+        self.conn.commit()
+
+    def get_top_messages_today(self, chat_id: int):
+        self.cursor.execute("""
+            SELECT user_id, messages_today FROM users
+            WHERE chat_id = ? AND messages_today > 0
+            ORDER BY messages_today DESC
+            """, (chat_id,))
+        return self.cursor.fetchall()
+
+    def get_top_messages_month(self, chat_id: int):
+        self.cursor.execute("""
+            SELECT user_id, messages_month FROM users
+            WHERE chat_id = ? AND messages_month > 0
+            ORDER BY messages_month DESC
+            """, (chat_id,))
+        return self.cursor.fetchall()
+
+    def get_top_silent(self, chat_id: int, threshold_time: int):
+        self.cursor.execute("""
+            SELECT user_id, last_message FROM users
+            WHERE chat_id = ? AND last_message < ?
+            ORDER BY last_message ASC
+            """, (chat_id, threshold_time))
+        return self.cursor.fetchall()
+
+    def get_top_warnings(self, chat_id: int):
+        self.cursor.execute("""
+            SELECT user_id, warnings FROM users
+            WHERE chat_id = ? AND warnings > 0
+            ORDER BY warnings DESC
+            """, (chat_id,))
+        return self.cursor.fetchall()
+
+    def reset_messages_today(self):
+        self.cursor.execute("UPDATE users SET messages_today = 0")
+        self.conn.commit()
+
+    def reset_messages_month(self):
+        self.cursor.execute("UPDATE users SET messages_month = 0")
+        self.conn.commit()
+
+
+DB = Database()