From d2706368629b9d39f9126660cd5f6724e4d12f88 Mon Sep 17 00:00:00 2001
From: 283375 <bilibili283375@163.com>
Date: Tue, 1 Oct 2024 00:34:22 +0800
Subject: [PATCH] fix(db.models): use `ArcaeaLanguage` enum for localized
 objects

---
 .../database/models/v5/arcaea.py              | 21 +++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/arcaea_offline/database/models/v5/arcaea.py b/src/arcaea_offline/database/models/v5/arcaea.py
index 67835b5..922db66 100644
--- a/src/arcaea_offline/database/models/v5/arcaea.py
+++ b/src/arcaea_offline/database/models/v5/arcaea.py
@@ -1,10 +1,14 @@
 from typing import List, Optional
 
-from sqlalchemy import ForeignKey, and_, func, select
+from sqlalchemy import Enum, ForeignKey, and_, func, select
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 from sqlalchemy_utils import create_view
 
-from arcaea_offline.constants.enums import ArcaeaRatingClass, ArcaeaSongSide
+from arcaea_offline.constants.enums.arcaea import (
+    ArcaeaLanguage,
+    ArcaeaRatingClass,
+    ArcaeaSongSide,
+)
 
 from .base import ModelsV5Base, ModelsV5ViewBase, ReprHelper
 
@@ -20,6 +24,11 @@ __all__ = [
 ]
 
 
+_ArcaeaLanguageEnumType = Enum(
+    ArcaeaLanguage, native_enum=False, values_callable=lambda e: [x.value for x in e]
+)
+
+
 class Pack(ModelsV5Base, ReprHelper):
     __tablename__ = "packs"
 
@@ -40,7 +49,7 @@ class PackLocalized(ModelsV5Base, ReprHelper):
     id: Mapped[str] = mapped_column(
         ForeignKey(Pack.id, onupdate="CASCADE", ondelete="NO ACTION")
     )
-    lang: Mapped[str]
+    lang: Mapped[ArcaeaLanguage] = mapped_column(_ArcaeaLanguageEnumType)
     name: Mapped[Optional[str]]
     description: Mapped[Optional[str]]
 
@@ -91,7 +100,7 @@ class SongLocalized(ModelsV5Base, ReprHelper):
     id: Mapped[str] = mapped_column(
         ForeignKey(Song.id, onupdate="CASCADE", ondelete="NO ACTION")
     )
-    lang: Mapped[str]
+    lang: Mapped[ArcaeaLanguage] = mapped_column(_ArcaeaLanguageEnumType)
     title: Mapped[Optional[str]]
     source: Mapped[Optional[str]]
 
@@ -107,7 +116,7 @@ class SongSearchWord(ModelsV5Base, ReprHelper):
     id: Mapped[str] = mapped_column(
         ForeignKey(Song.id, onupdate="CASCADE", ondelete="NO ACTION")
     )
-    lang: Mapped[str]
+    lang: Mapped[ArcaeaLanguage] = mapped_column(_ArcaeaLanguageEnumType)
     type: Mapped[int] = mapped_column(comment="1: title, 2: artist")
     value: Mapped[str]
 
@@ -163,7 +172,7 @@ class DifficultyLocalized(ModelsV5Base, ReprHelper):
     rating_class: Mapped[ArcaeaRatingClass] = mapped_column(
         ForeignKey(Difficulty.rating_class, onupdate="CASCADE", ondelete="NO ACTION")
     )
-    lang: Mapped[str]
+    lang: Mapped[ArcaeaLanguage] = mapped_column(_ArcaeaLanguageEnumType)
     title: Mapped[Optional[str]]
     artist: Mapped[Optional[str]]