diff --git a/main.py b/main.py index 0c8fdeb..9478238 100644 --- a/main.py +++ b/main.py @@ -1,18 +1,19 @@ import json -from functools import wraps import logging import os import re import sqlite3 from dataclasses import dataclass from datetime import datetime +from functools import wraps from pprint import pprint -from dotenv import load_dotenv -from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup, ParseMode import telegram -from telegram.ext import (CallbackContext, CommandHandler, MessageHandler, - Updater) +from dotenv import load_dotenv +from telegram import (InlineKeyboardButton, InlineKeyboardMarkup, ParseMode, + Update) +from telegram.ext import (CallbackContext, CallbackQueryHandler, + CommandHandler, MessageHandler, Updater) from models import Link @@ -59,6 +60,7 @@ class Bot: self.dispatcher.add_handler(MessageHandler( filters=None, callback=self.message)) + self.dispatcher.add_handler(CallbackQueryHandler(self.update)) @property def db(self): @@ -134,6 +136,30 @@ class Bot: context.bot.send_message( update.effective_chat.id, f"Added {self._natural_count(len(links), 'link', 'links')} to your list.") + @private_only + def update(self, update: Update, context: CallbackContext): + user_id = update.effective_user.id + action, link_id = update.callback_query.data.split(":", 1) + with self.db as db: + l = Link.get(db, user_id, link_id) + if l is None: + context.bot.send_message( + user_id, "Couldn't find the link you were looking for.") + return + + if action == "mark_as_read": + l.mark_as_read(db) + context.bot.send_message( + user_id, "Marked 1 link as read.") + update.callback_query.message.delete() + elif action == "delete": + l.delete(db) + update.callback_query.message.delete() + context.bot.send_message( + user_id, "Deleted 1 link.") + + + def run(self): self.updater.start_polling() diff --git a/models.py b/models.py index 80fca93..a50e699 100644 --- a/models.py +++ b/models.py @@ -1,6 +1,7 @@ import sqlite3 from dataclasses import dataclass from datetime import datetime +from os import link from typing import Dict, List, Tuple @@ -58,6 +59,21 @@ class Link: assert int(r[0][0]) self.id = r[0][0] + def mark_as_read(self, db: sqlite3.Connection): + assert self.id is not None + self.read_at = datetime.now() + cur = db.cursor() + cur.execute( + "UPDATE links SET read_at=? WHERE user_id=? AND id=?", (self.read_at, self.user_id, self.id)) + + def delete(self, db: sqlite3.Connection): + assert self.id is not None + cur = db.cursor() + cur.execute( + "DELETE FROM links WHERE user_id=? AND id=?", (self.user_id, self.id)) + self.id = None + + @classmethod def _get(cls, db: sqlite3.Connection, where: str, values: Tuple = ()) -> List['Link']: cur = db.cursor() @@ -67,3 +83,11 @@ class Link: @classmethod def get_unread(cls, db: sqlite3.Connection, user_id: str) -> List['Link']: return cls._get(db, "user_id = ? AND read_at IS NULL", (user_id,)) + + @classmethod + def get(cls, db: sqlite3.Connection, user_id: str, link_id: int) -> 'Link': + rows = list(cls._get(db, "user_id = ? AND id = ?", (user_id, link_id))) + if len(rows) < 1: + return None + else: + return rows[0]