Add helper function to load the right Environment based on file exten…

textworld/envs/__init__.py

@@ -1,7 +1,19 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT license.
+import re
 
 from textworld.envs.glulx.git_glulx import GitGlulxEnv
 from textworld.envs.zmachine.jericho import JerichoEnv
 from textworld.envs.tw import TextWorldEnv
 from textworld.envs.wrappers.tw_inform7 import TWInform7
+
+
+def _guess_backend(path):
+    # Guess the backend from the extension.
+    if path.endswith(".ulx"):
+        return GitGlulxEnv
+    elif re.search(r"\.z[1-8]", path):
+        return JerichoEnv
+
+    msg = "Unsupported game format: {}".format(path)
+    raise ValueError(msg)

textworld/helpers.py

@@ -8,9 +8,7 @@
 from textworld.core import EnvInfos, Environment, Agent
 from textworld.generator.game import Game, GameOptions
 
-from textworld.envs import GitGlulxEnv
-from textworld.envs import JerichoEnv
-from textworld.envs import TextWorldEnv
+import textworld.envs
 from textworld.envs import TWInform7
 
 from textworld.agents import HumanAgent
@@ -41,23 +39,10 @@ def start(path: str, request_infos: Optional[EnvInfos] = None,
         raise IOError(msg)
 
     # Guess the backend from the extension.
-    backend = "zmachine"
-    if path.endswith(".ulx"):
-        backend = "glulx"
-    elif path.endswith(".json"):
-        backend = "json"
-
-    if backend == "zmachine":
-        env = JerichoEnv(request_infos)
-    elif backend == "glulx":
-        env = GitGlulxEnv(request_infos)
-    elif backend == "json":
-        env = TextWorldEnv(request_infos)
-    else:
-        msg = "Unsupported backend: {}".format(backend)
-        raise ValueError(msg)
-
-    if TWInform7.compatible(path) and backend != "json":
+    Env = textworld.envs._guess_backend(path)
+    env = Env(request_infos)
+
+    if TWInform7.compatible(path):
         wrappers = [TWInform7] + list(wrappers)
 
     # Apply all wrappers