Переход между бранчами, vkbottle

from tortoise import Tortoise

​

from vkbottle import Bot, Message

from vkbottle.branch import ClsBranch, rule_disposal

from vkbottle.framework.framework.branch.database_branch import DatabaseBranch

from vkbottle.rule import VBMLRule

from tortoise_models import UserState

from vkbottle.keyboard import Keyboard, Text

from vkbottle.rule import AbstractMessageRule

​

import json

​

"""

Database branch manager itself

"""

​

class ComeTo(AbstractMessageRule):

    async def check(self, message: Message):

        # Функция check вызывается при проверке правила

        if message.payload and "to_branch" in message.payload:

            return True

​

class MysqlBranch(DatabaseBranch):

    async def set_user(self, uid: int, branch: str, context: str):

        """This method should make user's state or update it if exists"""

        user = await UserState.get_or_none(uid=uid)

        if user is not None:

            user.branch = branch

            user.context = context

            return await user.save()

        await UserState.create(uid=uid, branch=branch, context=context)

​

    async def get_user(self, uid: int):

        """This method should return a tuple of two strings: branch name of the user and context"""

        user = await UserState.get(uid=uid)

        return user.branch, user.context

​

    async def all_users(self):

        """This method should return user_ids of all stated users"""

        return [user.uid async for user in UserState.all()]

​

    async def delete_user(self, uid: int):

        pass

​

​

bot = Bot('TOKEN')

bot.branch = MysqlBranch()

​

class ChildBranch1(ClsBranch):

    @rule_disposal(VBMLRule("/котнекст <word>"))

    async def say(self, ans: Message, word: str):

        self.context["word"] = word

        return f"В моем контексте слово: {word}"

​

    async def branch(self, ans: Message, *args):

        keyboard = Keyboard(one_time=False)

        keyboard.add_row()

        keyboard.add_button(Text(label="Главное меню", payload={"to_branch": "main_branch"}), color="primary")

        user = await bot.api.users.get(ans.from_id)

        await ans(f"{user[0].first_name}, ты сейчас в Дочернем меню 1!", keyboard=keyboard.generate())

​

class ChildBranch2(ClsBranch):

    @rule_disposal(VBMLRule("/котнекст <word>"))

    async def say(self, ans: Message, word: str):

        self.context["word"] = word

        return f"В моем контексте слово: {word}"

​

    async def branch(self, ans: Message, *args):

        keyboard = Keyboard(one_time=False)

        keyboard.add_row()

        keyboard.add_button(Text(label="Главное меню", payload={"to_branch": "main_branch"}), color="primary")

        user = await bot.api.users.get(ans.from_id)

        await ans(f"{user[0].first_name}, ты сейчас в Дочернем меню 2!", keyboard=keyboard.generate())

​

class MainBranch(ClsBranch):

    @rule_disposal(ComeTo())

    async def come_to(self, ans: Message, *args):

        #await bot.branch.add(ans.peer_id, json.loads(ans.payload['to_branch']))

        await bot.branch.add(ans.peer_id, "child_branch_1") # Не меняет текущий бранч пользователю

​

    @rule_disposal(VBMLRule("/котнекст <word>"))

    async def say(self, ans: Message, word: str):

        self.context["word"] = word

        return f"В моем контексте слово: {word}"

​

    async def branch(self, ans: Message, *args):

        keyboard = Keyboard(one_time=False)

        keyboard.add_row()

        keyboard.add_button(Text(label="Дочернее меню 1", payload={"to_branch": "child_branch_1"}), color="primary")

        keyboard.add_row()

        keyboard.add_button(Text(label="Дочернее меню 2", payload={"to_branch": "child_branch_2"}), color="primary")

        user = await bot.api.users.get(ans.from_id)

        await ans(f"{user[0].first_name}, ты сейчас в Главном меню!", keyboard=keyboard.generate())

​

@bot.on.message_handler(commands=["начать"] )

async def start(ans: Message):

    await bot.branch.add(ans.peer_id, "main_branch")

    return "Ты попадаешь в Главный бранч!"

​

​

async def init_db():

    await Tortoise.init(

        db_url="sqlite://users.db", modules={"models": ["tortoise_models"]}

    )

    await Tortoise.generate_schemas()

​

​

bot.branch.add_branch(MainBranch, "main_branch")

bot.branch.add_branch(ChildBranch1, "child_branch_1")

bot.branch.add_branch(ChildBranch2, "child_branch_2")

bot.run_polling(on_startup=init_db)