From df543440770d97a526ffaf8fef26aa23aa934ba6 Mon Sep 17 00:00:00 2001
From: Acly <aclysia@gmail.com>
Date: Wed, 24 Jul 2024 12:00:07 +0200
Subject: [PATCH] Translate: parse language directives included in the text

---
 translation.py | 52 ++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 42 insertions(+), 10 deletions(-)

diff --git a/translation.py b/translation.py
index bece0bb..bd028be 100644
--- a/translation.py
+++ b/translation.py
@@ -1,7 +1,14 @@
-"""Text translation using Argos Translate."""
+"""Text translation using Argos Translate.
+
+The node takes text input and translates it to English. The text may contain any
+number of language directives in the form `lang:xx` where `xx` is a two-letter
+language code. Text fragments after a language directives are translated.
+If the language is `en` text is passed through unmodified.
+"""
 
 import re
 from functools import cache
+from typing import NamedTuple
 
 
 @cache
@@ -51,20 +58,45 @@ def translate(text: str, language: str):
 class Translate:
     @staticmethod
     def INPUT_TYPES():
-        return {
-            "required": {
-                "text": ("STRING", {"multiline": True}),
-                "language": ([code for code, _ in available_languages()],),
-                "target": (["en"],),
-            }
-        }
+        return {"required": {"text": ("STRING", {"multiline": True})}}
 
     CATEGORY = "external_tooling"
     RETURN_TYPES = ("STRING",)
     FUNCTION = "translate"
 
-    def translate(self, text: str, language: str, target: str):
-        return (translate(text, language),)
+    def translate(self, text: str):
+        chunks = Chunk.parse(text)
+        print(chunks)
+        translated = " ".join(translate(c.text, c.lang) for c in chunks)
+        return (translated,)
+
+
+_lang_regex = re.compile(r"(lang:\w\w)")
+
+
+class Chunk(NamedTuple):
+    text: str
+    lang: str
+
+    @staticmethod
+    def parse(text: str):
+        languages = [code for code, name in available_languages()] + ["en"]
+        chunks: list[Chunk] = []
+        lang = "en"
+        last = 0
+        for m in _lang_regex.finditer(text):
+            if m.start() > 0:
+                chunks.append(Chunk(text[last : m.start()].strip(), lang))
+            last = m.end()
+            lang = m.group(0)[5:]
+            if lang not in languages:
+                raise ValueError(
+                    f"Invalid language directive {m.group(0)} - {lang} is not a known language code."
+                    f" Available languages: {', '.join(languages)}"
+                )
+        if last < len(text):
+            chunks.append(Chunk(text[last:].strip(), lang))
+        return [c for c in chunks if c.text != ""]
 
 
 _embedding_regex = re.compile(r"(embedding:[^\s,]+)")