from pyrogram.types import InlineKeyboardMarkup


class InlineKeyboard(InlineKeyboardMarkup):
    def __init__(self, row_width=3):
        self.inline_keyboard = list()
        super().__init__(inline_keyboard=self.inline_keyboard)
        self.row_width = row_width

    def add(self, *args):
        self.inline_keyboard = [
            args[i:i + self.row_width]
            for i in range(0, len(args), self.row_width)
        ]

    def row(self, *args):
        self.inline_keyboard.append([button for button in args])
