From 05a29b6ee655e963c42363bccf7d8ba44e59779f Mon Sep 17 00:00:00 2001
From: argoneus <argoneuscze@gmail.com>
Date: Thu, 20 Feb 2020 23:57:09 +0100
Subject: [PATCH] translate command now works with global languages

---
 diva/quest.py      |  2 --
 utage/translate.py | 61 ++++++++++++++++++++++++++++++++++------------
 2 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/diva/quest.py b/diva/quest.py
index 5a6d27d..f2342a1 100644
--- a/diva/quest.py
+++ b/diva/quest.py
@@ -18,8 +18,6 @@ def scene_mst(qdb, old_path, dir_in, languages):
 
     locations = load_locations(dir_in)
 
-    print(locations)
-
     for id, scene in scenes.items():
         i = 0
         while True:
diff --git a/utage/translate.py b/utage/translate.py
index 8b70398..6ea96b4 100644
--- a/utage/translate.py
+++ b/utage/translate.py
@@ -4,6 +4,7 @@ import io
 import json
 import os
 import traceback
+from contextlib import suppress
 
 from functools import partial
 from glob import glob
@@ -49,9 +50,17 @@ def translate_file(file_in, tsv_out_dir, json_out_dir, is_global):
         tsv_output_path = os.path.join(
             tsv_out_dir, event_folder, "Scenario", "{}_t.tsv".format(id_num)
         )
-        json_output_path = os.path.join(
-            json_out_dir, event_folder, "{}_translations_jpn.json".format(id_num),
-        )
+        languages = ["jpn"]
+        if is_global:
+            languages += ["enm", "zh", "ko"]
+        json_output_paths = {
+            lang: os.path.join(
+                json_out_dir,
+                event_folder,
+                "{}_translations_{}.json".format(id_num, lang),
+            )
+            for lang in languages
+        }
 
         # need to create output paths and avoid races when threading
         if not os.path.exists(os.path.dirname(tsv_output_path)):
@@ -61,16 +70,21 @@ def translate_file(file_in, tsv_out_dir, json_out_dir, is_global):
                 if e.errno != errno.EEXIST:
                     raise
 
-        if not os.path.exists(os.path.dirname(json_output_path)):
+        if not os.path.exists(os.path.dirname(json_output_paths["jpn"])):
             try:
-                os.makedirs(os.path.dirname(json_output_path))
+                os.makedirs(os.path.dirname(json_output_paths["jpn"]))
             except OSError as e:
                 if e.errno != errno.EEXIST:
                     raise
 
         with open(file_in, "r") as tsv_file:
             tsv = csv.DictReader(tsv_file, delimiter="\t", quotechar='"')
-            tsv_keyed, json_str = process_tsv(tsv, id_num)
+
+            t_fieldnames = tsv.fieldnames
+            if is_global:
+                t_fieldnames = tsv.fieldnames[: tsv.fieldnames.index("English") + 1]
+
+            tsv_keyed, json_str = process_tsv(tsv, id_num, is_global)
 
             # csv handles newlines, don't set it in io.open
             with io.open(tsv_output_path, "w", newline="") as tsv_out:
@@ -79,28 +93,39 @@ def translate_file(file_in, tsv_out_dir, json_out_dir, is_global):
                     delimiter="\t",
                     quotechar='"',
                     lineterminator="\n",
-                    fieldnames=tsv.fieldnames,
+                    fieldnames=t_fieldnames,
                     extrasaction="ignore",
                 )
                 writer.writeheader()
                 for row in tsv_keyed:
                     writer.writerow(row)
 
-        with io.open(json_output_path, "w", newline="\n") as json_out:
-            # we don't want to sort these because they're _1, ..., _10, etc
-            json.dump(
-                json_str, json_out, ensure_ascii=False, indent="\t", sort_keys=False,
-            )
+        for lang in languages:
+            if not json_str[lang]:
+                continue
+            with io.open(json_output_paths[lang], "w", newline="\n") as json_out:
+                # we don't want to sort these because they're _1, ..., _10, etc
+                json.dump(
+                    json_str[lang],
+                    json_out,
+                    ensure_ascii=False,
+                    indent="\t",
+                    sort_keys=False,
+                )
 
     except Exception:
         traceback.print_exc()
         print("Error processing file: " + file_in)
 
 
-def process_tsv(tsv, id_num):
+def process_tsv(tsv, id_num, is_global):
     i = 0
     tsv_keyed = []
-    key_dict = {}
+    key_dict = {"jpn": {}}
+    if is_global:
+        key_dict["enm"] = {}
+        key_dict["zh"] = {}
+        key_dict["ko"] = {}
 
     for row in tsv:
         try:
@@ -110,8 +135,14 @@ def process_tsv(tsv, id_num):
             if row["Text"] and row["Text"].strip():
                 key = "{}_{}".format(id_num, i)
                 i += 1
+                key_dict["jpn"][key] = row["Text"]
+                if is_global:
+                    with suppress(KeyError):
+                        if row["English"]:
+                            key_dict["enm"][key] = row["English"]
+                        key_dict["zh"][key] = row["Chinese"]
+                        key_dict["ko"][key] = row["Korean"]
                 row["English"] = key
-                key_dict[key] = row["Text"]
             tsv_keyed.append(row)
         except Exception:
             traceback.print_exc()