code clean up

fedn/cli/run_cmd.py

@@ -43,14 +43,12 @@ def check_yaml_exists(path):
         click.echo(f"Could not find fedn.yaml in {path}")
         exit(-1)
     return yaml_file
-def delete_virtual_environment(keep_venv, dispatcher):
-    if not keep_venv:
-        # delete the virtualenv
-        if dispatcher.python_env_path:
-            logger.info(f"Removing virtualenv {dispatcher.python_env_path}")
-            shutil.rmtree(dispatcher.python_env_path)
-        else:
-            logger.warning("No virtualenv found to remove.")
+def delete_virtual_environment(dispatcher):
+    if dispatcher.python_env_path:
+        logger.info(f"Removing virtualenv {dispatcher.python_env_path}")
+        shutil.rmtree(dispatcher.python_env_path)
+    else:
+        logger.warning("No virtualenv found to remove.")
 @main.group("run")
 @click.pass_context
 def run_cmd(ctx):
@@ -61,7 +59,7 @@ def run_cmd(ctx):
 @click.option("-p", "--path", required=True, help="Path to package directory containing fedn.yaml")
 @click.option("-i", "--input", required=True, help="Path to input model" )
 @click.option("-o", "--output", required=True, help="Path to write the output JSON containing validation metrics")
-@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set to False doesn't remove venv")
+@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set doesn't remove venv")
 @click.pass_context
 def validate_cmd(ctx, path, input, output, keep_venv):
     """Execute 'validate' entrypoint in fedn.yaml.
@@ -82,12 +80,13 @@ def validate_cmd(ctx, path, input, output, keep_venv):
     dispatcher = Dispatcher(config, path)
     _ = dispatcher._get_or_create_python_env()
     dispatcher.run_cmd("validate {} {}".format(input, output))
-    delete_virtual_environment(keep_venv, dispatcher)
+    if not keep_venv:
+        delete_virtual_environment(dispatcher)
 @run_cmd.command("train")
 @click.option("-p", "--path", required=True, help="Path to package directory containing fedn.yaml")
 @click.option("-i", "--input", required=True, help="Path to input model parameters" )
 @click.option("-o", "--output", required=True, help="Path to write the updated model parameters ")
-@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set to False doesn't remove venv")
+@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set doesn't remove venv")
 @click.pass_context
 def train_cmd(ctx, path, input, output, keep_venv):
     """Execute 'train' entrypoint in fedn.yaml.
@@ -108,11 +107,11 @@ def train_cmd(ctx, path, input, output, keep_venv):
     dispatcher = Dispatcher(config, path)
     _ = dispatcher._get_or_create_python_env()
     dispatcher.run_cmd("train {} {}".format(input, output))
-    delete_virtual_environment(keep_venv, dispatcher)
-
+    if not keep_venv:
+        delete_virtual_environment(dispatcher)
 @run_cmd.command("startup")
 @click.option("-p", "--path", required=True, help="Path to package directory containing fedn.yaml")
-@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set to False doesn't remove venv")
+@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set doesn't remove venv")
 @click.pass_context
 def startup_cmd(ctx, path, keep_venv):
     """Execute 'startup' entrypoint in fedn.yaml.
@@ -132,12 +131,12 @@ def startup_cmd(ctx, path, keep_venv):
     dispatcher = Dispatcher(config, path)
     _ = dispatcher._get_or_create_python_env()
     dispatcher.run_cmd("startup")
-    delete_virtual_environment(keep_venv, dispatcher)
-
+    if not keep_venv:
+        delete_virtual_environment(dispatcher)
 
 @run_cmd.command("build")
 @click.option("-p", "--path", required=True, help="Path to package directory containing fedn.yaml")
-@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set to False doesn't remove venv")
+@click.option("-v", "--keep-venv", is_flag=True, required=False, help="flag if set doesn't remove venv")
 @click.pass_context
 def build_cmd(ctx, path, keep_venv):
     """Execute 'build' entrypoint in fedn.yaml.
@@ -158,8 +157,8 @@ def build_cmd(ctx, path, keep_venv):
     dispatcher = Dispatcher(config, path)
     _ = dispatcher._get_or_create_python_env()
     dispatcher.run_cmd("build")
-    delete_virtual_environment(keep_venv, dispatcher)
-
+    if not keep_venv:
+        delete_virtual_environment(dispatcher)
 
 
 @run_cmd.command("client")