flashbot/words_progress_db.py at main · jbenbee/flashbot · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import os.path
import threading
from typing import Optional

import pandas as pd

from item import Item


class WordsProgressDB:
    def __init__(self, db_path):
        self.db_path = db_path
        if not os.path.exists(self.db_path):
            self.progress_df = pd.DataFrame(columns=['chat_id', 'word_id', 'num_reps', 'e_factor', 'last_interval', 'last_review_date', 'next_review_date', 'to_ignore'])
        self.progress_df = pd.read_csv(self.db_path)
        self.progress_df['last_review_date'] = pd.to_datetime(self.progress_df['last_review_date'])
        self.progress_df['next_review_date'] = pd.to_datetime(self.progress_df['next_review_date']).dt.date
        self.progress_df['word_id'] = self.progress_df['word_id'].astype(int)
        self.progress_df['chat_id'] = self.progress_df['chat_id'].astype(int)
        self._lock = threading.Lock()

    def get_progress_df(self):
        self._lock.acquire()
        pdf_cpy = self.progress_df.copy()
        self._lock.release()
        return pdf_cpy

    def save_progress(self):
        self._lock.acquire()
        self.progress_df.to_csv(self.db_path, index=False)
        self._lock.release()

    def add_word_to_progress(self, chat_id, word_id):
        self._lock.acquire()
        new_row = pd.DataFrame([{'chat_id': chat_id, 'word_id': word_id, 'num_reps': 0.0, 'to_ignore': False,
             'e_factor': 2.5, 'last_interval': 0, 'last_review_date': None, 'next_review_date': None}])[self.progress_df.columns].astype(self.progress_df.dtypes)
        self.progress_df = pd.concat([self.progress_df, new_row], ignore_index=True)
        self._lock.release()

    def ignore_word(self, chat_id, word_id):
        self._lock.acquire()
        chat_word_progress = self.progress_df[(self.progress_df['chat_id'] == chat_id) &
                                              (self.progress_df['word_id'] == word_id)]
        if chat_word_progress.shape[0] == 0:
            self.add_word_to_progress(chat_id, word_id)

        mask = (self.progress_df['chat_id'] == chat_id) & (self.progress_df['word_id'] == word_id)
        if self.progress_df[mask].shape[0] != 1:
            self._lock.release()
            raise ValueError('Number of rows satisfying the condition must be exactly 1.')
        self.progress_df.loc[mask, 'to_ignore'] = True
        self._lock.release()

    def get_word_progress(self, chat_id, word_id) -> Optional[Item]:
        self._lock.acquire()
        chat_word_progress = self.progress_df[(self.progress_df['chat_id'] == chat_id) &
                                              (self.progress_df['word_id'] == word_id)]
        if chat_word_progress.shape[0] == 0:
            self._lock.release()
            return None

        mask = (self.progress_df['chat_id'] == chat_id) & (self.progress_df['word_id'] == word_id)
        row = self.progress_df[mask]
        if row.shape[0] != 1:
            self._lock.release()
            raise ValueError('Number of rows satisfying the condition must be exactly 1.')
        item = Item(e_factor=row['e_factor'].item(), num_reps=row['num_reps'].item(),
                    next_review_date=row['next_review_date'].item(), last_review_date=row['last_review_date'].item(),
                    last_interval=row['last_interval'].item(), word_id=word_id)
        self._lock.release()
        return item

    def set_word_progress(self, chat_id, word_id, item) -> None:
        self._lock.acquire()
        mask = (self.progress_df['chat_id'] == chat_id) & (self.progress_df['word_id'] == word_id)

        chat_word_progress = self.progress_df[mask]
        if chat_word_progress.shape[0] != 1:
            self._lock.release()
            raise ValueError(f'Number of rows satisfying the condition must be exactly 1, found {chat_word_progress.shape[0]}.')
        self.progress_df.loc[mask, 'num_reps'] = item.num_reps
        self.progress_df.loc[mask, 'e_factor'] = item.e_factor
        self.progress_df.loc[mask, 'next_review_date'] = item.next_review_date
        self.progress_df.loc[mask, 'last_review_date'] = item.last_review_date
        self.progress_df.loc[mask, 'last_interval'] = item.last_interval
        self._lock.release()

    def remove_progress(self, chat_id, word_ids):
        self._lock.acquire()
        mask = (self.progress_df['chat_id'] == chat_id) & (self.progress_df['word_id'].isin(word_ids))
        self.progress_df = self.progress_df[~mask]
        self._lock.release()

    def release_lock(self):
        if self._lock.locked():
            self._lock.release()