make i18n more extendable

ricequant · Feb 15, 2017 · 58fce45 · 58fce45
1 parent 8f80a2a
commit 58fce45
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 26 deletions.
diff --git a/rqalpha/__main__.py b/rqalpha/__main__.py
@@ -28,7 +28,7 @@
 
 from .cache_control import set_cache_policy, CachePolicy
 from .utils.click_helper import Date
-from .utils.i18n import set_locale
+from .utils.i18n import localization
 from .utils.config import parse_config
 
 
@@ -119,7 +119,7 @@ def run(**kwargs):
         locale.setlocale(locale.LC_ALL, "en_US.UTF-8")
         locale.setlocale(locale.LC_CTYPE, "en_US.UTF-8")
         os.environ['TZ'] = 'Asia/Shanghai'
-        set_locale(["zh_Hans_CN"])
+        localization.set_locale(["zh_Hans_CN"])
     except Exception as e:
         if os.name != 'nt':
             raise

diff --git a/rqalpha/utils/i18n.py b/rqalpha/utils/i18n.py
@@ -18,30 +18,35 @@
 from gettext import NullTranslations, translation
 from .logger import system_log
 
-translation_dir = os.path.join(
-    os.path.dirname(
-        os.path.abspath(
-            __file__,
-        ),
-    ),
-    "translations"
-)
-current_translation = NullTranslations()
-
-
-def set_locale(locales, trans_dir=translation_dir):
-    global current_translation
-
-    try:
-        current_translation = translation(
-            domain="messages",
-            localedir=trans_dir,
-            languages=locales,
-        )
-    except Exception as e:
-        system_log.debug(e)
-        current_translation = NullTranslations()
+
+class Localization(object):
+
+    def __init__(self, trans=None):
+        self.trans = NullTranslations() if trans is None else trans
+
+    def set_locale(self, locales, trans_dir=None):
+        try:
+            if trans_dir is None:
+                trans_dir = os.path.join(
+                    os.path.dirname(
+                        os.path.abspath(
+                            __file__,
+                        ),
+                    ),
+                    "translations"
+                )
+            self.trans = translation(
+                domain="messages",
+                localedir=trans_dir,
+                languages=locales,
+            )
+        except Exception as e:
+            system_log.debug(e)
+            self.trans = NullTranslations()
+
+
+localization = Localization()
 
 
 def gettext(message):
-    return current_translation.gettext(message)
+    return localization.trans.gettext(message)