Пол пользователя

2026-02-04 13:00:20 +03:00
parent 05c4773a50
commit d51bbe04ba
6 changed files with 154 additions and 48 deletions
--- a/src/database.py
+++ b/src/database.py
@@ -13,6 +13,7 @@ class UserState(Enum):
    """Состояния пользовательской сессии."""

    INTRO = "intro"
+    SPECIFY_GENDER = "specify_gender"
    NO_MORE_SCENARIOS = "no_more_scenarios"
    FIRST_REPLICA = "first_replica"
    SHOW_REPLICA = "show_replica"
@@ -32,6 +33,7 @@ class User:
    id: int  # dataset_speaker_id
    telegram_id: int
    created_at: datetime
+    gender: str | None  # "male" или "female"


@dataclass
@@ -51,6 +53,7 @@ class Replica:
    speaker_id: int  # в рамках сценария
    replica_index: int  # порядок в сценарии (0-indexed)
    text: str
+    gender: str  # "male" или "female"


@dataclass
@@ -101,7 +104,8 @@ def init_db() -> None:
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                telegram_id INTEGER UNIQUE NOT NULL,
-                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                gender TEXT
            );

            CREATE TABLE IF NOT EXISTS scenarios (
@@ -115,6 +119,7 @@ def init_db() -> None:
                speaker_id INTEGER NOT NULL,
                replica_index INTEGER NOT NULL,
                text TEXT NOT NULL,
+                gender TEXT NOT NULL,
                FOREIGN KEY (scenario_id) REFERENCES scenarios(id),
                UNIQUE(scenario_id, replica_index)
            );
@@ -152,14 +157,6 @@ def init_db() -> None:
            CREATE INDEX IF NOT EXISTS idx_recordings_scenario 
                ON recordings(scenario_id);
        """)
-
-        # Миграция: добавляем колонку duration если её нет
-        cursor = conn.execute("PRAGMA table_info(recordings)")
-        columns = [row[1] for row in cursor.fetchall()]
-        if "duration" not in columns:
-            conn.execute("ALTER TABLE recordings ADD COLUMN duration REAL DEFAULT 0.0")
-            logger.info("Добавлена колонка duration в таблицу recordings")
-
        conn.commit()

    logger.info("База данных инициализирована")
@@ -172,7 +169,8 @@ def get_or_create_user(telegram_id: int) -> User:
    """Получает или создаёт пользователя по telegram_id."""
    with get_connection() as conn:
        cursor = conn.execute(
-            "SELECT id, telegram_id, created_at FROM users WHERE telegram_id = ?",
+            "SELECT id, telegram_id, created_at, gender FROM users "
+            "WHERE telegram_id = ?",
            (telegram_id,),
        )
        row = cursor.fetchone()
@@ -182,18 +180,22 @@ def get_or_create_user(telegram_id: int) -> User:
                id=row["id"],
                telegram_id=row["telegram_id"],
                created_at=row["created_at"],
+                gender=row["gender"],
            )

        cursor = conn.execute(
            "INSERT INTO users (telegram_id) VALUES (?) "
-            "RETURNING id, telegram_id, created_at",
+            "RETURNING id, telegram_id, created_at, gender",
            (telegram_id,),
        )
        row = cursor.fetchone()
        conn.commit()
        logger.info(f"Создан новый пользователь: dataset_speaker_id={row['id']}")
        return User(
-            id=row["id"], telegram_id=row["telegram_id"], created_at=row["created_at"]
+            id=row["id"],
+            telegram_id=row["telegram_id"],
+            created_at=row["created_at"],
+            gender=row["gender"],
        )


@@ -201,7 +203,8 @@ def get_user_by_telegram_id(telegram_id: int) -> User | None:
    """Получает пользователя по telegram_id."""
    with get_connection() as conn:
        cursor = conn.execute(
-            "SELECT id, telegram_id, created_at FROM users WHERE telegram_id = ?",
+            "SELECT id, telegram_id, created_at, gender FROM users "
+            "WHERE telegram_id = ?",
            (telegram_id,),
        )
        row = cursor.fetchone()
@@ -210,10 +213,22 @@ def get_user_by_telegram_id(telegram_id: int) -> User | None:
                id=row["id"],
                telegram_id=row["telegram_id"],
                created_at=row["created_at"],
+                gender=row["gender"],
            )
        return None


+def update_user_gender(user_id: int, gender: str) -> None:
+    """Обновляет пол пользователя."""
+    with get_connection() as conn:
+        conn.execute(
+            "UPDATE users SET gender = ? WHERE id = ?",
+            (gender, user_id),
+        )
+        conn.commit()
+        logger.info(f"Обновлён пол пользователя {user_id}: {gender}")
+
+
 # === Scenarios CRUD ===


@@ -258,15 +273,18 @@ def get_all_scenarios() -> list[Scenario]:
 # === Replicas CRUD ===


-def create_replicas(scenario_id: str, replicas: list[tuple[int, int, str]]) -> None:
-    """Создаёт реплики. replicas: [(speaker_id, replica_index, text), ...]"""
+def create_replicas(
+    scenario_id: str, replicas: list[tuple[int, int, str, str]]
+) -> None:
+    """Создаёт реплики. replicas: [(speaker_id, replica_index, text, gender), ...]"""
    with get_connection() as conn:
        conn.executemany(
-            "INSERT INTO replicas (scenario_id, speaker_id, replica_index, text) "
-            "VALUES (?, ?, ?, ?)",
+            "INSERT INTO replicas "
+            "(scenario_id, speaker_id, replica_index, text, gender) "
+            "VALUES (?, ?, ?, ?, ?)",
            [
-                (scenario_id, speaker_id, idx, text)
-                for speaker_id, idx, text in replicas
+                (scenario_id, speaker_id, idx, text, gender)
+                for speaker_id, idx, text, gender in replicas
            ],
        )
        conn.commit()
@@ -276,8 +294,8 @@ def get_replicas_for_scenario(scenario_id: str) -> list[Replica]:
    """Получает все реплики сценария."""
    with get_connection() as conn:
        cursor = conn.execute(
-            "SELECT id, scenario_id, speaker_id, replica_index, text FROM replicas "
-            "WHERE scenario_id = ? ORDER BY replica_index",
+            "SELECT id, scenario_id, speaker_id, replica_index, text, gender "
+            "FROM replicas WHERE scenario_id = ? ORDER BY replica_index",
            (scenario_id,),
        )
        return [
@@ -287,6 +305,7 @@ def get_replicas_for_scenario(scenario_id: str) -> list[Replica]:
                speaker_id=row["speaker_id"],
                replica_index=row["replica_index"],
                text=row["text"],
+                gender=row["gender"],
            )
            for row in cursor.fetchall()
        ]
@@ -296,8 +315,9 @@ def get_replicas_for_track(scenario_id: str, speaker_id: int) -> list[Replica]:
    """Получает реплики для конкретной дорожки (speaker_id в сценарии)."""
    with get_connection() as conn:
        cursor = conn.execute(
-            "SELECT id, scenario_id, speaker_id, replica_index, text FROM replicas "
-            "WHERE scenario_id = ? AND speaker_id = ? ORDER BY replica_index",
+            "SELECT id, scenario_id, speaker_id, replica_index, text, gender "
+            "FROM replicas WHERE scenario_id = ? AND speaker_id = ? "
+            "ORDER BY replica_index",
            (scenario_id, speaker_id),
        )
        return [
@@ -307,6 +327,7 @@ def get_replicas_for_track(scenario_id: str, speaker_id: int) -> list[Replica]:
                speaker_id=row["speaker_id"],
                replica_index=row["replica_index"],
                text=row["text"],
+                gender=row["gender"],
            )
            for row in cursor.fetchall()
        ]
@@ -323,6 +344,16 @@ def get_track_speaker_ids(scenario_id: str) -> list[int]:
        return [row["speaker_id"] for row in cursor.fetchall()]


+def get_scenario_genders(scenario_id: str) -> set[str]:
+    """Получает список полов, представленных в сценарии."""
+    with get_connection() as conn:
+        cursor = conn.execute(
+            "SELECT DISTINCT gender FROM replicas WHERE scenario_id = ?",
+            (scenario_id,),
+        )
+        return {row["gender"] for row in cursor.fetchall()}
+
+
 # === Recordings CRUD ===