From abeee7e38548c7e6193ceccba9900e4dac7ee4ac Mon Sep 17 00:00:00 2001
From: Philippe-Cholet <phcholet@orange.fr>
Date: Sun, 28 Mar 2021 16:19:19 +0200
Subject: [PATCH] Add option --load-only language codes

The user will not have to install language packages from "argos-translate" that he/she has no use for.
---
 app/app.py  |  2 +-
 app/init.py | 25 ++++++++++++++++++++++---
 app/main.py |  4 ++++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/app/app.py b/app/app.py
index d673ff2..8ac8e7c 100644
--- a/app/app.py
+++ b/app/app.py
@@ -49,7 +49,7 @@ def get_routes_limits(default_req_limit, api_keys_db):
 def create_app(args):
     if not args.offline:
         from app.init import boot
-        boot()
+        boot(args.load_only)
 
     from app.language import languages
     app = Flask(__name__)
diff --git a/app/init.py b/app/init.py
index 9cf83f7..4724738 100644
--- a/app/init.py
+++ b/app/init.py
@@ -5,11 +5,11 @@ import os, glob, shutil, zipfile
 from app.language import languages
 import polyglot
 
-def boot():
-    check_and_install_models()
+def boot(load_only=None):
+    check_and_install_models(load_only_lang_codes=load_only)
     check_and_install_transliteration()
 
-def check_and_install_models(force=False):
+def check_and_install_models(force=False, load_only_lang_codes=None):
     if len(package.get_installed_packages()) < 2 or force:
         # Update package definitions from remote
         print("Updating language models")
@@ -19,6 +19,25 @@ def check_and_install_models(force=False):
         available_packages = package.load_available_packages()
         print("Found %s models" % len(available_packages))
 
+        if load_only_lang_codes is not None:
+            # load_only_lang_codes: List[str] (codes)
+            # Ensure the user does not use any unavailable language code.
+            unavailable_lang_codes = set(load_only_lang_codes)
+            for pack in available_packages:
+                unavailable_lang_codes -= {pack.from_code, pack.to_code}
+            if unavailable_lang_codes:
+                raise ValueError('Unavailable language codes: %s.' % ','.join(sorted(unavailable_lang_codes)))
+            # Keep only the packages that have both from_code and to_code in our list.
+            available_packages = [
+                pack
+                for pack in available_packages
+                if pack.from_code in load_only_lang_codes
+                and pack.to_code in load_only_lang_codes
+            ]
+            if not available_packages:
+                raise ValueError('no available package')
+            print("Keep %s models" % len(available_packages))
+
         # Download and install all available packages
         for available_package in available_packages:
             print("Downloading %s (%s) ..." % (available_package, available_package.package_version))
diff --git a/app/main.py b/app/main.py
index ffc3ef2..af60e0f 100644
--- a/app/main.py
+++ b/app/main.py
@@ -1,4 +1,5 @@
 import argparse
+import operator
 from app.app import create_app
 
 def main():
@@ -29,6 +30,9 @@ def main():
                         help="Use offline")
     parser.add_argument('--api-keys', default=False, action="store_true",
                         help="Enable API keys database for per-user rate limits lookup")
+    parser.add_argument('--load-only', type=operator.methodcaller('split', ','),
+                        metavar='<comma-separated language codes>',
+                        help='Set available languages (ar,de,en,es,fr,ga,hi,it,ja,ko,pt,ru,zh)')
     
 
     args = parser.parse_args()