add choice

aiidateam · Jan 21, 2022 · 8e4b70a · 8e4b70a
1 parent 0cfae86
commit 8e4b70a
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 7 deletions.
diff --git a/aiida/cmdline/commands/cmd_config.py b/aiida/cmdline/commands/cmd_config.py
@@ -17,7 +17,7 @@
 from aiida.cmdline.commands.cmd_verdi import verdi
 from aiida.cmdline.params import arguments
 from aiida.cmdline.utils import echo
-from aiida.manage.configuration import downgrade_config, get_config_path
+from aiida.manage.configuration import MIGRATIONS, downgrade_config, get_config_path
 from aiida.manage.configuration.settings import DEFAULT_CONFIG_INDENT_SIZE
 
 
@@ -197,12 +197,12 @@ def verdi_config_caching(disabled):
 
 
 @verdi_config.command('downgrade')
-@click.argument('version', type=int)
+@click.argument('version', type=click.Choice({str(m.down_revision) for m in MIGRATIONS}))
 def verdi_config_downgrade(version):
     """Print a configuration, downgraded to a specific version."""
     path = Path(get_config_path())
     echo.echo_report(f'Downgrading configuration to v{version}: {path}')
     config = json.loads(path.read_text(encoding='utf8'))
-    downgrade_config(config, version)
+    downgrade_config(config, int(version))
     path.write_text(json.dumps(config, indent=DEFAULT_CONFIG_INDENT_SIZE), encoding='utf8')
     echo.echo_success('Downgraded')
diff --git a/aiida/manage/__init__.py b/aiida/manage/__init__.py
@@ -39,6 +39,7 @@
     'ConfigValidationError',
     'DEFAULT_DBINFO',
     'DeliveryFailed',
+    'MIGRATIONS',
     'OLDEST_COMPATIBLE_CONFIG_VERSION',
     'Option',
     'Postgres',

diff --git a/aiida/manage/configuration/__init__.py b/aiida/manage/configuration/__init__.py
@@ -23,6 +23,7 @@
     'CURRENT_CONFIG_VERSION',
     'Config',
     'ConfigValidationError',
+    'MIGRATIONS',
     'OLDEST_COMPATIBLE_CONFIG_VERSION',
     'Option',
     'Profile',

diff --git a/aiida/manage/configuration/migrations/__init__.py b/aiida/manage/configuration/migrations/__init__.py
@@ -18,6 +18,7 @@
 
 __all__ = (
     'CURRENT_CONFIG_VERSION',
+    'MIGRATIONS',
     'OLDEST_COMPATIBLE_CONFIG_VERSION',
     'check_and_migrate_config',
     'config_needs_migrating',

diff --git a/aiida/manage/configuration/migrations/migrations.py b/aiida/manage/configuration/migrations/migrations.py
@@ -14,7 +14,7 @@
 
 __all__ = (
     'CURRENT_CONFIG_VERSION', 'OLDEST_COMPATIBLE_CONFIG_VERSION', 'get_current_version', 'check_and_migrate_config',
-    'config_needs_migrating', 'upgrade_config', 'downgrade_config'
+    'config_needs_migrating', 'upgrade_config', 'downgrade_config', 'MIGRATIONS'
 )
 
 ConfigType = Dict[str, Any]
@@ -196,7 +196,7 @@ def downgrade(self, config: ConfigType) -> None:
                 config['options'][current] = config['options'].pop(new)
 
 
-_MIGRATIONS = (Initial, AddProfileUuid, SimplifyDefaultProfiles, AddMessageBroker, SimplifyOptions)
+MIGRATIONS = (Initial, AddProfileUuid, SimplifyDefaultProfiles, AddMessageBroker, SimplifyOptions)
 
 
 def get_current_version(config):
@@ -218,7 +218,7 @@ def get_oldest_compatible_version(config):
 def upgrade_config(
     config: ConfigType,
     target: int = CURRENT_CONFIG_VERSION,
-    migrations: Iterable[Type[SingleMigration]] = _MIGRATIONS
+    migrations: Iterable[Type[SingleMigration]] = MIGRATIONS
 ) -> ConfigType:
     """Run the registered configuration migrations up to the target version.
 
@@ -246,7 +246,7 @@ def upgrade_config(
 
 
 def downgrade_config(
-    config: ConfigType, target: int, migrations: Iterable[Type[SingleMigration]] = _MIGRATIONS
+    config: ConfigType, target: int, migrations: Iterable[Type[SingleMigration]] = MIGRATIONS
 ) -> ConfigType:
     """Run the registered configuration migrations down to the target version.