From b183c8290d55cb5aaad8c7ca218eef337983a6b4 Mon Sep 17 00:00:00 2001
From: Tohru <65994850+Tohrusky@users.noreply.github.com>
Date: Mon, 19 Aug 2024 23:00:21 +0800
Subject: [PATCH] feat: preset ass config (#5)

---
 pyproject.toml    |  2 +-
 tests/test_sub.py |  7 ++++-
 yuisub/prompt.py  |  2 +-
 yuisub/sub.py     | 66 ++++++++++++++++++++++++++++++++---------------
 4 files changed, 53 insertions(+), 24 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 522e878..ef809cd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,7 @@ license = "GPL-3.0-only"
 name = "yuisub"
 readme = "README.md"
 repository = "https://github.com/TensoRaws/yuisub"
-version = "0.0.4"
+version = "0.0.5"
 
 # Requirements
 [tool.poetry.dependencies]
diff --git a/tests/test_sub.py b/tests/test_sub.py
index 04bc644..fcbb313 100644
--- a/tests/test_sub.py
+++ b/tests/test_sub.py
@@ -19,9 +19,14 @@ def test_audio() -> None:
     sub.save(util.projectPATH / "assets" / "test.audio.ass")
 
 
-@pytest.mark.skipif(os.environ.get("GITHUB_ACTIONS") == "true", reason="Skipping test when running on CI")
 def test_bilingual() -> None:
     sub = load(util.TEST_ENG_SRT)
+    _ = bilingual(sub, sub)
+
+
+@pytest.mark.skipif(os.environ.get("GITHUB_ACTIONS") == "true", reason="Skipping test when running on CI")
+def test_bilingual_2() -> None:
+    sub = load(util.TEST_ENG_SRT)
 
     sub_zh = translate(
         sub=sub,
diff --git a/yuisub/prompt.py b/yuisub/prompt.py
index 532f3a3..9b64748 100644
--- a/yuisub/prompt.py
+++ b/yuisub/prompt.py
@@ -20,7 +20,7 @@ def anime_prompt(bangumi_url: str | None = None) -> str:
     return (
         """
             你的目标是把动漫新番的台词翻译成中文，要翻译得自然、流畅和地道，使用贴合二次元的表达方式。
-            请注意，由于字幕可能是通过 AI 生成的，所以会有一些错误的地方，请你在翻译时尽量保持逻辑性。
+            字幕可能是通过 AI 生成的，请你在翻译时尽量保持逻辑性和连贯性。
             此外，我可能会给你一些动漫相关的信息。请注意，当人名等专有名词出现时，严格按照我提供的信息进行翻译。
 
             """
diff --git a/yuisub/sub.py b/yuisub/sub.py
index 6765d22..45f86bd 100644
--- a/yuisub/sub.py
+++ b/yuisub/sub.py
@@ -10,6 +10,27 @@
 from yuisub.llm import Translator
 from yuisub.prompt import ORIGIN
 
+PRESET_STYLES: dict[str, SSAStyle] = {
+    "zh": SSAStyle(
+        alignment=Alignment.BOTTOM_CENTER,
+        primarycolor=Color(215, 215, 215),
+        fontsize=18,
+        fontname="Microsoft YaHei",
+        bold=True,
+        shadow=0,
+        outline=1,
+        outlinecolor=Color(198, 107, 107),
+    ),
+    "origin": SSAStyle(
+        alignment=Alignment.BOTTOM_CENTER,
+        primarycolor=pysubs2.Color(249, 246, 240),
+        fontsize=10,
+        fontname="Microsoft YaHei",
+        shadow=0,
+        outline=0.5,
+    ),
+}
+
 
 def load(sub_path: Path | str, encoding: str = "utf-8") -> SSAFile:
     """
@@ -24,7 +45,14 @@ def load(sub_path: Path | str, encoding: str = "utf-8") -> SSAFile:
 
 
 @retry(wait=wait_random(min=3, max=5), stop=stop_after_attempt(5))
-def translate(sub: SSAFile, model: str, api_key: str, base_url: str, bangumi_url: str | None = None) -> SSAFile:
+def translate(
+    sub: SSAFile,
+    model: str,
+    api_key: str,
+    base_url: str,
+    bangumi_url: str | None = None,
+    styles: dict[str, SSAStyle] | None = None,
+) -> SSAFile:
     """
     Translate subtitle file to Chinese
 
@@ -33,6 +61,7 @@ def translate(sub: SSAFile, model: str, api_key: str, base_url: str, bangumi_url
     :param api_key: llm api_key
     :param base_url: llm base_url
     :param bangumi_url: anime bangumi url
+    :param styles: subtitle styles, default is PRESET_STYLES
     :return:
     """
 
@@ -57,43 +86,38 @@ async def _wait_tasks() -> None:
     asyncio.run(_wait_tasks())
 
     # generate Chinese subtitle
+    if styles is None:
+        styles = PRESET_STYLES
+
     sub_zh = deepcopy(sub)
+    sub_zh.styles = styles
     for i, _ in enumerate(sub):
         sub_zh[i].text = trans_list[i]
+        sub_zh[i].style = "zh"
 
     return sub_zh
 
 
-def bilingual(sub_origin: SSAFile, sub_zh: SSAFile) -> SSAFile:
+def bilingual(
+    sub_origin: SSAFile,
+    sub_zh: SSAFile,
+    styles: dict[str, SSAStyle] | None = None,
+) -> SSAFile:
     """
     Generate bilingual subtitle file
 
     :param sub_origin: Origin subtitle
     :param sub_zh: Chinese subtitle
+    :param styles: subtitle styles, default is PRESET_STYLES
     :return:
     """
 
     # generate bilingual subtitle
+    if styles is None:
+        styles = PRESET_STYLES
+
     sub_bilingual = SSAFile()
-    sub_bilingual.styles = {
-        "zh": SSAStyle(
-            alignment=Alignment.BOTTOM_CENTER,
-            primarycolor=Color(255, 192, 203),
-            fontsize=16,
-            fontname="Microsoft YaHei",
-            bold=True,
-            shadow=0,
-            outline=0.2,
-        ),
-        "origin": SSAStyle(
-            alignment=Alignment.BOTTOM_CENTER,
-            primarycolor=pysubs2.Color(249, 246, 240),
-            fontsize=12,
-            fontname="Microsoft YaHei",
-            shadow=0,
-            outline=0.5,
-        ),
-    }
+    sub_bilingual.styles = styles
 
     for e in sub_origin:
         e.style = "origin"