rapnote/rhyme_backend.py at master · codebdbd/rapnote · 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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import sqlite3
import os
from pathlib import Path

# Используем путь к базе данных в директории проекта
DB_FILE = str(Path(__file__).parent / "db" / "rhymes.db")


def _escape_like(pattern: str) -> str:
    """Escape LIKE metacharacters so they are treated as literals."""
    return pattern.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")


class RhymeEngine:
    STRESS_CHAR = '\u0301'
    VOWELS = "аеёиоуыэюяіїє"

    @staticmethod
    def get_stress_position(word_with_stress: str) -> int:
        if 'ё' in word_with_stress.lower():
             v_cnt = 0
             for char in word_with_stress.lower()[::-1]:
                 if char in RhymeEngine.VOWELS:
                     if char == 'ё': return v_cnt
                     v_cnt += 1

        if RhymeEngine.STRESS_CHAR not in word_with_stress:
            return -1

        w_rev = word_with_stress[::-1]
        vowel_count = 0
        stress_found = False

        for char in w_rev:
            if char == RhymeEngine.STRESS_CHAR:
                stress_found = True
                continue
            if char.lower() in RhymeEngine.VOWELS:
                if stress_found: return vowel_count
                vowel_count += 1

        return -1

    @staticmethod
    def find_rhymes(word_input: str, lang: str) -> list[tuple[str, int]]:
        if not os.path.exists(DB_FILE):
            return []

        target_clean = word_input.lower().strip().replace(RhymeEngine.STRESS_CHAR, "")
        if not target_clean: return []

        if len(target_clean) < 2:
            ending = target_clean
        else:
            ending = target_clean[-2:]

        target_stress_pos = RhymeEngine.get_stress_position(word_input)

        with sqlite3.connect(DB_FILE) as conn:
            c = conn.cursor()

            if target_stress_pos == -1:
                try:
                    c.execute("SELECT word FROM words WHERE clean = ? AND lang = ? LIMIT 1", (target_clean, lang))
                    row = c.fetchone()
                    if row:
                        target_stress_pos = RhymeEngine.get_stress_position(row[0])
                except sqlite3.Error:
                    pass

            escaped_ending = _escape_like(ending)
            query = """
                SELECT word, clean FROM words
                WHERE lang = ?
                AND clean LIKE ? ESCAPE '\\'
            """
            c.execute(query, (lang, f"%{escaped_ending}"))
            results = c.fetchall()

        candidates = []

        for row in results:
            original_word = row[0]
            clean_word = row[1]

            # --- ФИЛЬТРЫ ---
            if clean_word == target_clean: continue
            if '-' in clean_word: continue

            # Фильтр Заглавных (Имена, Города)
            if original_word and original_word[0].isupper():
                continue

            score = 0

            # === НОВАЯ БАЛАНСИРОВКА ОЧКОВ ===

            # 1. БУКВЫ
            if clean_word.endswith(ending):
                score = 50 # База (Зеленый минимум)

                # Бонус за 3 буквы (+20)
                # Если совпало: 50 + 20 = 70
                if len(target_clean) >= 3 and clean_word.endswith(target_clean[-3:]):
                    score += 20

            # 2. РИТМ
            if target_stress_pos != -1:
                cand_stress = RhymeEngine.get_stress_position(original_word)

                if cand_stress != -1:
                    if cand_stress == target_stress_pos:
                        # Бонус за Ритм (+20)

                        # СЦЕНАРИЙ А (Простая рифма "Чат"):
                        # 50 (буквы) + 20 (ритм) = 70 (ЗЕЛЕНЫЙ)

                        # СЦЕНАРИЙ Б (Крутая рифма "Результат"):
                        # 50 (буквы) + 20 (3 буквы) + 20 (ритм) = 90 (СИНИЙ)
                        score += 20
                    else:
                        # Штраф за сбитый ритм
                        score -= 20

            # Показываем всё, что набрало хотя бы 30 (чтобы список был полным)
            if score >= 30:
                candidates.append((original_word, score))

        # Сортировка: Очки -> Длина
        candidates.sort(key=lambda x: (-x[1], len(x[0])))

        return candidates[:2000]