From 355b53c48518617662c3df3a923c3f31951fc0b3 Mon Sep 17 00:00:00 2001
From: Shabaz Patel <shabaz.b.patel@gmail.com>
Date: Mon, 4 Jun 2018 01:32:34 -0700
Subject: [PATCH] adding env setup in cli command

---
 datmo/cli/command/environment.py            | 14 ++++++++++++++
 datmo/cli/command/tests/test_environment.py | 15 +++++++++++++++
 datmo/cli/parser.py                         | 11 +++++++++++
 datmo/core/util/lang/en.py                  |  4 ++++
 4 files changed, 44 insertions(+)

diff --git a/datmo/cli/command/environment.py b/datmo/cli/command/environment.py
index a72d6996..b55c565b 100644
--- a/datmo/cli/command/environment.py
+++ b/datmo/cli/command/environment.py
@@ -19,6 +19,20 @@ def environment(self):
         self.parse(["--help"])
         return True
 
+    def setup(self, **kwargs):
+        libraries = kwargs.get("libraries", None)
+        available_libraries = self.environment_controller.get_current_libraries()
+        if not libraries:
+            for index, library in enumerate(available_libraries):
+                self.cli_helper.echo("(%s) %s" % (index+1, library))
+            libraries = self.cli_helper.prompt(
+                __("prompt", "cli.environment.setup.libraries"))
+        if libraries in available_libraries:
+            options = {"libraries": libraries}
+            return self.environment_controller.setup(options=options)
+        else:
+            self.cli_helper.echo(__("error", "cli.environment.setup.argument", libraries))
+
     def create(self, **kwargs):
         self.cli_helper.echo(__("info", "cli.environment.create"))
         environment_obj = self.environment_controller.create(kwargs)
diff --git a/datmo/cli/command/tests/test_environment.py b/datmo/cli/command/tests/test_environment.py
index 1efab231..620bf076 100644
--- a/datmo/cli/command/tests/test_environment.py
+++ b/datmo/cli/command/tests/test_environment.py
@@ -53,6 +53,21 @@ def __set_variables(self):
         self.environment_command = EnvironmentCommand(self.temp_dir,
                                                       self.cli_helper)
 
+    def test_environment_setup(self):
+        # Setup the environement by passing libraries
+        self.__set_variables()
+        datmo_environment_folder = os.path.join(self.temp_dir,
+                                                "datmo_environment")
+
+        definition_filepath = os.path.join(datmo_environment_folder,
+                                           "Dockerfile")
+        test_libraries = 'xgboost:cpu'
+        self.environment_command.parse(["environment", "setup", "--libraries", test_libraries])
+        result = self.environment_command.execute()
+
+        assert result and os.path.isfile(definition_filepath) and \
+               "datmo" in open(definition_filepath, "r").read()
+
     def test_environment_create(self):
         # 1) Environment definition file in `datmo_environment` folder
         # 2) Environment definition file passed as an option
diff --git a/datmo/cli/parser.py b/datmo/cli/parser.py
index bc22f8d8..1d8537b5 100644
--- a/datmo/cli/parser.py
+++ b/datmo/cli/parser.py
@@ -52,6 +52,17 @@ def get_datmo_parser():
     environment_subcommand_parsers = environment_parser.add_subparsers(
         title="subcommands", dest="subcommand")
 
+    environment_setup = environment_subcommand_parsers.add_parser(
+        "setup", help="setup environment")
+    environment_setup.add_argument(
+        "--libraries",
+        dest="libraries",
+        default=None,
+        type=str,
+        help=
+        "libraries/framework to be used for environment (e.g. xgboost:cpu)"
+    )
+
     environment_create = environment_subcommand_parsers.add_parser(
         "create", help="create environment")
     environment_create.add_argument(
diff --git a/datmo/core/util/lang/en.py b/datmo/core/util/lang/en.py
index aa76d035..13975b48 100644
--- a/datmo/core/util/lang/en.py
+++ b/datmo/core/util/lang/en.py
@@ -37,6 +37,8 @@
             "{foo} - {bar}",
         "cli.general.tuple.test":
             "%s, %s",
+        "cli.environment.setup.argument":
+            "No library exists with this name: %s",
         "cli.environment.create":
             "Creating a new environment",
         "cli.environment.create.success":
@@ -328,6 +330,8 @@
     "trace": {},
     "fatal": {},
     "prompt": {
+        "cli.environment.setup.libraries":
+            "Please choose one of the above images",
         "cli.project.init.name":
             "Enter name for the project",
         "cli.project.init.description":