Use a wrapper function to determine the available environments

Author: filwaitman

divio_cli/cli.py

@@ -21,6 +21,7 @@
     open_project_cloud_site,
     print_package_renamed_warning,
     table,
+    get_available_environments,
 )
 from .validators.addon import validate_addon
 from .validators.boilerplate import validate_boilerplate
@@ -208,7 +209,7 @@ def sort_projects(items):
 @click.option(
     "--backup/--no-backup", default=None, help="Take a backup on deployment."
 )
-@click.argument("stage", default="test", type=click.Choice(['test', 'live']))
+@click.argument("stage", default="test", type=click.Choice(get_available_environments()))
 @allow_remote_id_override
 @click.pass_obj
 def project_deploy(obj, remote_id, stage, backup):
@@ -217,7 +218,7 @@ def project_deploy(obj, remote_id, stage, backup):
 
 
 @project.command(name="deploy-log")
-@click.argument("stage", default="test", type=click.Choice(['test', 'live']))
+@click.argument("stage", default="test", type=click.Choice(get_available_environments()))
 @allow_remote_id_override
 @click.pass_obj
 def project_deploy_log(obj, remote_id, stage):
@@ -418,7 +419,7 @@ def project_pull():
 
 
 @project_pull.command(name="db")
-@click.argument("stage", default="test", type=click.Choice(['test', 'live']))
+@click.argument("stage", default="test", type=click.Choice(get_available_environments()))
 @allow_remote_id_override
 @click.pass_obj
 def pull_db(obj, remote_id, stage):
@@ -432,7 +433,7 @@ def pull_db(obj, remote_id, stage):
 
 
 @project_pull.command(name="media")
-@click.argument("stage", default="test", type=click.Choice(['test', 'live']))
+@click.argument("stage", default="test", type=click.Choice(get_available_environments()))
 @allow_remote_id_override
 @click.pass_obj
 def pull_media(obj, remote_id, stage):
@@ -449,7 +450,7 @@ def project_push():
 
 
 @project_push.command(name="db")
-@click.argument("stage", default="test", type=click.Choice(['test', 'live']))
+@click.argument("stage", default="test", type=click.Choice(get_available_environments()))
 @click.option(
     "-d",
     "--dumpfile",
@@ -487,7 +488,7 @@ def push_db(obj, remote_id, stage, dumpfile, noinput):
 
 
 @project_push.command(name="media")
-@click.argument("stage", default="test", type=click.Choice(['test', 'live']))
+@click.argument("stage", default="test", type=click.Choice(get_available_environments()))
 @click.option(
     "--noinput", is_flag=True, default=False, help="Don't ask for confirmation"
 )

divio_cli/utils.py

@@ -319,6 +319,10 @@ def json_dumps_unicode(d, **kwargs):
     return json.dumps(d, ensure_ascii=False, **kwargs).encode("utf-8")
 
 
+def get_available_environments():
+    return ['test', 'live']
+
+
 class Map(dict):
     """
     A dictionary which also allows accessing values by dot notation.