From 8d5afedb26fa4fc81d359a74c2418754fbb65024 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Thu, 12 Jun 2025 14:36:11 +0200 Subject: [PATCH 1/6] wip --- src/aiida/cmdline/commands/cmd_computer.py | 62 ++++++++- src/aiida/cmdline/params/options/config.py | 1 + src/aiida/cmdline/params/options/main.py | 18 +++ src/aiida/cmdline/utils/template_config.py | 150 +++++++++++++++++++++ 4 files changed, 225 insertions(+), 6 deletions(-) create mode 100644 src/aiida/cmdline/utils/template_config.py diff --git a/src/aiida/cmdline/commands/cmd_computer.py b/src/aiida/cmdline/commands/cmd_computer.py index 67123f8a24..8512fa2e35 100644 --- a/src/aiida/cmdline/commands/cmd_computer.py +++ b/src/aiida/cmdline/commands/cmd_computer.py @@ -8,6 +8,7 @@ ########################################################################### """`verdi computer` command.""" +import json import pathlib import traceback from copy import deepcopy @@ -22,6 +23,7 @@ from aiida.cmdline.utils import echo, echo_tabulate from aiida.cmdline.utils.common import validate_output_filename from aiida.cmdline.utils.decorators import with_dbenv +from aiida.cmdline.utils.template_config import load_and_process_template from aiida.common.exceptions import EntryPointError, ValidationError from aiida.plugins.entry_point import get_entry_point_names @@ -270,6 +272,19 @@ def set_computer_builder(ctx, param, value): return value +# Helper function to set template vars in context +def set_template_vars_in_context(ctx, param, value): + """Set template variables in the context for the config provider to use.""" + if value: + try: + template_var_dict = json.loads(value) + # Store template vars in context for the config provider + ctx._template_vars = template_var_dict + except json.JSONDecodeError as e: + raise click.BadParameter(f'Invalid JSON in template-vars: {e}') + return value + +# Modified computer setup command @verdi_computer.command('setup') @options_computer.LABEL() @options_computer.HOSTNAME() @@ -284,23 +299,57 @@ def set_computer_builder(ctx, param, value): @options_computer.USE_DOUBLE_QUOTES() @options_computer.PREPEND_TEXT() @options_computer.APPEND_TEXT() -@options.NON_INTERACTIVE() -@options.CONFIG_FILE() +# @options.NON_INTERACTIVE() +# @options.CONFIG_FILE() # Keep the original config option for backward compatibility +@options.TEMPLATE_FILE() # Add our new template option +@options.TEMPLATE_VARS() @click.pass_context @with_dbenv() -def computer_setup(ctx, non_interactive, **kwargs): +def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): """Create a new computer.""" from aiida.orm.utils.builders.computer import ComputerBuilder + print("HELLO") - if kwargs['label'] in get_computer_names(): + # Handle template variables + template_var_dict = None + if template_vars: + try: + template_var_dict = json.loads(template_vars) + except json.JSONDecodeError as e: + echo.echo_critical(f'Invalid JSON in template-vars: {e}') + + # Process template if provided + if template: + try: + # Load and process the template + config_data = load_and_process_template( + template, + interactive=not non_interactive, + template_vars=template_var_dict + ) + + # Update kwargs with config file data + # Only update if the value wasn't explicitly provided on command line + for key, value in config_data.items(): + if key not in kwargs or kwargs[key] is None: + kwargs[key] = value + + except Exception as e: + echo.echo_critical(f'Error processing template: {e}') + + # Check for existing computer + if kwargs.get('label') and kwargs['label'] in get_computer_names(): echo.echo_critical( 'A computer called {c} already exists. ' 'Use "verdi computer duplicate {c}" to set up a new ' 'computer starting from the settings of {c}.'.format(c=kwargs['label']) ) - kwargs['transport'] = kwargs['transport'].name - kwargs['scheduler'] = kwargs['scheduler'].name + # Convert entry points to their names + if kwargs.get('transport'): + kwargs['transport'] = kwargs['transport'].name + if kwargs.get('scheduler'): + kwargs['scheduler'] = kwargs['scheduler'].name computer_builder = ComputerBuilder(**kwargs) try: @@ -321,6 +370,7 @@ def computer_setup(ctx, non_interactive, **kwargs): echo.echo_report(f' verdi -p {profile.name} computer configure {computer.transport_type} {computer.label}') + @verdi_computer.command('duplicate') @arguments.COMPUTER(callback=set_computer_builder) @options_computer.LABEL(contextual_default=partial(get_parameter_default, 'label')) diff --git a/src/aiida/cmdline/params/options/config.py b/src/aiida/cmdline/params/options/config.py index 99794b862d..1a3d74706d 100644 --- a/src/aiida/cmdline/params/options/config.py +++ b/src/aiida/cmdline/params/options/config.py @@ -75,6 +75,7 @@ def configuration_callback( except Exception as exception: raise click.BadOptionUsage(option_name, f'Error reading configuration file: {exception}', ctx) + import ipdb; ipdb.set_trace() valid_params = [param.name for param in ctx.command.params if param.name != option_name] specified_params = list(config.keys()) unknown_params = set(specified_params).difference(set(valid_params)) diff --git a/src/aiida/cmdline/params/options/main.py b/src/aiida/cmdline/params/options/main.py index 906e5168f7..dad0070a47 100644 --- a/src/aiida/cmdline/params/options/main.py +++ b/src/aiida/cmdline/params/options/main.py @@ -113,6 +113,8 @@ 'SORT', 'START_DATE', 'SYMLINK_CALCS', + 'TEMPLATE_FILE', + 'TEMPLATE_VARS', 'TIMEOUT', 'TRAJECTORY_INDEX', 'TRANSPORT', @@ -917,3 +919,19 @@ def set_log_level(ctx, _param, value): show_default=True, help='End date for node mtime range selection for node collection dumping.', ) + +# Add a new option for template variables in non-interactive mode +TEMPLATE_VARS = OverridableOption( + '--template-vars', + type=click.STRING, + help='JSON string containing template variable values for non-interactive mode. ' + 'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'', +) + +# Create an enhanced config file option +TEMPLATE_FILE = click.option( + '--template', + type=click.STRING, + help='Load computer setup from configuration file in YAML format (local path or URL). ' + 'Supports Jinja2 templates with interactive prompting.', +) diff --git a/src/aiida/cmdline/utils/template_config.py b/src/aiida/cmdline/utils/template_config.py new file mode 100644 index 0000000000..eb64b1bb78 --- /dev/null +++ b/src/aiida/cmdline/utils/template_config.py @@ -0,0 +1,150 @@ +from typing import Any, Dict, List, Optional + +import click +import requests +import yaml +from jinja2 import BaseLoader, Environment, meta + +from aiida.cmdline.utils import echo + + +class StringTemplateLoader(BaseLoader): + """Jinja2 loader that loads templates from strings.""" + + def __init__(self, template_string: str): + self.template_string = template_string + + def get_source(self, environment, template): + return self.template_string, None, lambda: True + + +def prompt_for_template_variables(template_variables: Dict[str, Any]) -> Dict[str, Any]: + """Prompt user for template variables based on metadata definitions.""" + values = {} + + echo.echo_report('Template variables detected. Please provide values:') + echo.echo('') + + for var_name, var_config in template_variables.items(): + key_display = var_config.get('key_display', var_name) + description = var_config.get('description', f'Value for {var_name}') + var_type = var_config.get('type', 'text') + default = var_config.get('default') + options = var_config.get('options', []) + + # Display help text + echo.echo(f'{click.style(key_display, fg="yellow")}') + echo.echo(f' {description}') + + if var_type == 'list' and options: + echo.echo(f' Options: {", ".join(options)}') + while True: + value = click.prompt(' Enter value', default=default, show_default=True if default else False) + if value in options: + values[var_name] = value + break + else: + echo.echo_error(f'Invalid option. Please choose from: {", ".join(options)}') + else: + value = click.prompt(' Enter value', default=default, show_default=True if default else False) + values[var_name] = value + + echo.echo('') + + return values + + +def detect_template_variables(template_content: str) -> List[str]: + """Detect Jinja2 variables in template content.""" + env = Environment(loader=StringTemplateLoader(template_content)) + ast = env.parse(template_content) + return list(meta.find_undeclared_variables(ast)) + + +def load_and_process_template( + file_path_or_url: str, interactive: bool = True, template_vars: Optional[Dict[str, Any]] = None +) -> Dict[str, Any]: + """Load and process a template configuration file.""" + + # Load content + if file_path_or_url.startswith(('http://', 'https://')): + try: + response = requests.get(file_path_or_url, timeout=10) + response.raise_for_status() + content = response.text + except requests.RequestException as e: + raise click.BadParameter(f'Failed to fetch URL {file_path_or_url}: {e}') + else: + try: + with open(file_path_or_url, 'r', encoding='utf-8') as f: + content = f.read() + except IOError as e: + raise click.BadParameter(f'Failed to read file {file_path_or_url}: {e}') + + # Parse YAML to get metadata + try: + full_config = yaml.safe_load(content) + except yaml.YAMLError as e: + raise click.BadParameter(f'Invalid YAML: {e}') + + # Extract metadata and template variables (if they exist) + metadata = full_config.pop('metadata', {}) + template_variables = metadata.get('template_variables', {}) + + # Detect variables that need values + detected_vars = detect_template_variables(content) + + # If no template variables detected, just return the config + if not detected_vars: + return full_config + + # Filter to only prompt for variables that are actually used and defined in metadata + vars_to_prompt = {var: config for var, config in template_variables.items() if var in detected_vars} + + if vars_to_prompt: + if interactive: + # Interactive prompting for template variables + template_values = prompt_for_template_variables(vars_to_prompt) + else: + # Non-interactive mode + if not template_vars: + raise click.BadParameter( + f'Template variables detected ({", ".join(detected_vars)}) but no values provided. ' + 'Use --template-vars to provide values in JSON format.' + ) + template_values = template_vars + + # Render the template with provided values + env = Environment(loader=StringTemplateLoader(content)) + template = env.from_string(content) + rendered_content = template.render(**template_values) + + # Parse the rendered YAML + try: + config = yaml.safe_load(rendered_content) + except yaml.YAMLError as e: + raise click.BadParameter(f'Invalid YAML after template rendering: {e}') + else: + # Template variables detected but none defined in metadata + # This could happen with simple Jinja variables like {{ username }} + if interactive: + echo.echo_warning(f'Template variables detected ({", ".join(detected_vars)}) but no metadata found.') + echo.echo_warning('You may need to provide values manually or the template may not render correctly.') + + if template_vars: + # Try to render with provided vars + env = Environment(loader=StringTemplateLoader(content)) + template = env.from_string(content) + rendered_content = template.render(**template_vars) + try: + config = yaml.safe_load(rendered_content) + except yaml.YAMLError as e: + raise click.BadParameter(f'Invalid YAML after template rendering: {e}') + else: + # Return original config and hope for the best + config = full_config + + # Remove metadata section if it exists + config.pop('metadata', None) + + return config From caec3747f6eeeffd4317004c37c23ed28433491c Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Wed, 18 Jun 2025 08:10:26 +0200 Subject: [PATCH 2/6] fix template_file overridable option --- src/aiida/cmdline/params/options/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aiida/cmdline/params/options/main.py b/src/aiida/cmdline/params/options/main.py index e970742c9b..7f98735fd0 100644 --- a/src/aiida/cmdline/params/options/main.py +++ b/src/aiida/cmdline/params/options/main.py @@ -922,7 +922,7 @@ def set_log_level(ctx, _param, value): ) # Create an enhanced config file option -TEMPLATE_FILE = click.option( +TEMPLATE_FILE = OverridableOption( '--template', type=click.STRING, help='Load computer setup from configuration file in YAML format (local path or URL). ' From 1887c0e1858c9cf0ce40df92c54382b37ca7e719 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 18 Jun 2025 06:12:29 +0000 Subject: [PATCH 3/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/aiida/cmdline/commands/cmd_computer.py | 9 ++++----- src/aiida/cmdline/params/options/config.py | 4 +++- src/aiida/cmdline/params/options/main.py | 4 ++-- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/aiida/cmdline/commands/cmd_computer.py b/src/aiida/cmdline/commands/cmd_computer.py index 8512fa2e35..e86b5d06fb 100644 --- a/src/aiida/cmdline/commands/cmd_computer.py +++ b/src/aiida/cmdline/commands/cmd_computer.py @@ -284,6 +284,7 @@ def set_template_vars_in_context(ctx, param, value): raise click.BadParameter(f'Invalid JSON in template-vars: {e}') return value + # Modified computer setup command @verdi_computer.command('setup') @options_computer.LABEL() @@ -308,7 +309,8 @@ def set_template_vars_in_context(ctx, param, value): def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): """Create a new computer.""" from aiida.orm.utils.builders.computer import ComputerBuilder - print("HELLO") + + print('HELLO') # Handle template variables template_var_dict = None @@ -323,9 +325,7 @@ def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): try: # Load and process the template config_data = load_and_process_template( - template, - interactive=not non_interactive, - template_vars=template_var_dict + template, interactive=not non_interactive, template_vars=template_var_dict ) # Update kwargs with config file data @@ -370,7 +370,6 @@ def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): echo.echo_report(f' verdi -p {profile.name} computer configure {computer.transport_type} {computer.label}') - @verdi_computer.command('duplicate') @arguments.COMPUTER(callback=set_computer_builder) @options_computer.LABEL(contextual_default=partial(get_parameter_default, 'label')) diff --git a/src/aiida/cmdline/params/options/config.py b/src/aiida/cmdline/params/options/config.py index 1a3d74706d..9e35a62250 100644 --- a/src/aiida/cmdline/params/options/config.py +++ b/src/aiida/cmdline/params/options/config.py @@ -75,7 +75,9 @@ def configuration_callback( except Exception as exception: raise click.BadOptionUsage(option_name, f'Error reading configuration file: {exception}', ctx) - import ipdb; ipdb.set_trace() + import ipdb + + ipdb.set_trace() valid_params = [param.name for param in ctx.command.params if param.name != option_name] specified_params = list(config.keys()) unknown_params = set(specified_params).difference(set(valid_params)) diff --git a/src/aiida/cmdline/params/options/main.py b/src/aiida/cmdline/params/options/main.py index 7f98735fd0..3db36d6adb 100644 --- a/src/aiida/cmdline/params/options/main.py +++ b/src/aiida/cmdline/params/options/main.py @@ -918,7 +918,7 @@ def set_log_level(ctx, _param, value): '--template-vars', type=click.STRING, help='JSON string containing template variable values for non-interactive mode. ' - 'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'', + 'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'', ) # Create an enhanced config file option @@ -926,5 +926,5 @@ def set_log_level(ctx, _param, value): '--template', type=click.STRING, help='Load computer setup from configuration file in YAML format (local path or URL). ' - 'Supports Jinja2 templates with interactive prompting.', + 'Supports Jinja2 templates with interactive prompting.', ) From 25488bb80814933a0476f31e64aba02cade4213f Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Wed, 18 Jun 2025 08:31:13 +0200 Subject: [PATCH 4/6] Placeholders are being prompted, but later still excepts. --- src/aiida/cmdline/commands/cmd_computer.py | 95 ++++++++++++++++++++- src/aiida/cmdline/params/options/config.py | 3 - src/aiida/cmdline/params/options/main.py | 98 ++++++++++++++++++++-- 3 files changed, 180 insertions(+), 16 deletions(-) diff --git a/src/aiida/cmdline/commands/cmd_computer.py b/src/aiida/cmdline/commands/cmd_computer.py index e86b5d06fb..4fca0d503f 100644 --- a/src/aiida/cmdline/commands/cmd_computer.py +++ b/src/aiida/cmdline/commands/cmd_computer.py @@ -285,6 +285,7 @@ def set_template_vars_in_context(ctx, param, value): return value + # Modified computer setup command @verdi_computer.command('setup') @options_computer.LABEL() @@ -300,9 +301,8 @@ def set_template_vars_in_context(ctx, param, value): @options_computer.USE_DOUBLE_QUOTES() @options_computer.PREPEND_TEXT() @options_computer.APPEND_TEXT() -# @options.NON_INTERACTIVE() -# @options.CONFIG_FILE() # Keep the original config option for backward compatibility -@options.TEMPLATE_FILE() # Add our new template option +@options.NON_INTERACTIVE() # Uncomment this line +@options.TEMPLATE_FILE() @options.TEMPLATE_VARS() @click.pass_context @with_dbenv() @@ -310,7 +310,10 @@ def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): """Create a new computer.""" from aiida.orm.utils.builders.computer import ComputerBuilder - print('HELLO') + # Debug output + print(f"Debug: non_interactive = {non_interactive}") + print(f"Debug: kwargs keys = {list(kwargs.keys())}") + print(f"Debug: ctx.default_map = {ctx.default_map}") # Handle template variables template_var_dict = None @@ -369,6 +372,90 @@ def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): profile = ctx.obj['profile'] echo.echo_report(f' verdi -p {profile.name} computer configure {computer.transport_type} {computer.label}') +# Modified computer setup command +# @verdi_computer.command('setup') +# @options_computer.LABEL() +# @options_computer.HOSTNAME() +# @options_computer.DESCRIPTION() +# @options_computer.TRANSPORT() +# @options_computer.SCHEDULER() +# @options_computer.SHEBANG() +# @options_computer.WORKDIR() +# @options_computer.MPI_RUN_COMMAND() +# @options_computer.MPI_PROCS_PER_MACHINE() +# @options_computer.DEFAULT_MEMORY_PER_MACHINE() +# @options_computer.USE_DOUBLE_QUOTES() +# @options_computer.PREPEND_TEXT() +# @options_computer.APPEND_TEXT() +# # @options.NON_INTERACTIVE() +# # @options.CONFIG_FILE() # Keep the original config option for backward compatibility +# @options.TEMPLATE_FILE() # Add our new template option +# @options.TEMPLATE_VARS() +# @click.pass_context +# @with_dbenv() +# def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): +# """Create a new computer.""" +# from aiida.orm.utils.builders.computer import ComputerBuilder + +# print('HELLO') + +# # Handle template variables +# template_var_dict = None +# if template_vars: +# try: +# template_var_dict = json.loads(template_vars) +# except json.JSONDecodeError as e: +# echo.echo_critical(f'Invalid JSON in template-vars: {e}') + +# # Process template if provided +# if template: +# try: +# # Load and process the template +# config_data = load_and_process_template( +# template, interactive=not non_interactive, template_vars=template_var_dict +# ) + +# # Update kwargs with config file data +# # Only update if the value wasn't explicitly provided on command line +# for key, value in config_data.items(): +# if key not in kwargs or kwargs[key] is None: +# kwargs[key] = value + +# except Exception as e: +# echo.echo_critical(f'Error processing template: {e}') + +# # Check for existing computer +# if kwargs.get('label') and kwargs['label'] in get_computer_names(): +# echo.echo_critical( +# 'A computer called {c} already exists. ' +# 'Use "verdi computer duplicate {c}" to set up a new ' +# 'computer starting from the settings of {c}.'.format(c=kwargs['label']) +# ) + +# # Convert entry points to their names +# if kwargs.get('transport'): +# kwargs['transport'] = kwargs['transport'].name +# if kwargs.get('scheduler'): +# kwargs['scheduler'] = kwargs['scheduler'].name + +# computer_builder = ComputerBuilder(**kwargs) +# try: +# computer = computer_builder.new() +# except (ComputerBuilder.ComputerValidationError, ValidationError) as e: +# echo.echo_critical(f'{type(e).__name__}: {e}') + +# try: +# computer.store() +# except ValidationError as err: +# echo.echo_critical(f'unable to store the computer: {err}. Exiting...') +# else: +# echo.echo_success(f'Computer<{computer.pk}> {computer.label} created') + +# echo.echo_report('Note: before the computer can be used, it has to be configured with the command:') + +# profile = ctx.obj['profile'] +# echo.echo_report(f' verdi -p {profile.name} computer configure {computer.transport_type} {computer.label}') + @verdi_computer.command('duplicate') @arguments.COMPUTER(callback=set_computer_builder) diff --git a/src/aiida/cmdline/params/options/config.py b/src/aiida/cmdline/params/options/config.py index 9e35a62250..99794b862d 100644 --- a/src/aiida/cmdline/params/options/config.py +++ b/src/aiida/cmdline/params/options/config.py @@ -75,9 +75,6 @@ def configuration_callback( except Exception as exception: raise click.BadOptionUsage(option_name, f'Error reading configuration file: {exception}', ctx) - import ipdb - - ipdb.set_trace() valid_params = [param.name for param in ctx.command.params if param.name != option_name] specified_params = list(config.keys()) unknown_params = set(specified_params).difference(set(valid_params)) diff --git a/src/aiida/cmdline/params/options/main.py b/src/aiida/cmdline/params/options/main.py index 3db36d6adb..7e58174262 100644 --- a/src/aiida/cmdline/params/options/main.py +++ b/src/aiida/cmdline/params/options/main.py @@ -913,18 +913,98 @@ def set_log_level(ctx, _param, value): help='End date for node mtime range selection for node collection dumping.', ) -# Add a new option for template variables in non-interactive mode -TEMPLATE_VARS = OverridableOption( - '--template-vars', - type=click.STRING, - help='JSON string containing template variable values for non-interactive mode. ' - 'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'', -) +import click + +from aiida.cmdline.utils import echo +from aiida.cmdline.utils.template_config import load_and_process_template + +from .overridable import OverridableOption + + +class TemplateOption(OverridableOption): + """Option that processes Jinja2 templates and updates Click context defaults.""" + + def __init__(self, *args, **kwargs): + """Initialize the template option.""" + kwargs.setdefault('is_eager', True) # Process template before other options + kwargs.setdefault('expose_value', False) # Don't pass template to the command function + super().__init__(*args, **kwargs) + + def __call__(self, **kwargs): + """Create the option with updated kwargs.""" + kw_copy = self.kwargs.copy() + kw_copy.update(kwargs) + + # Set the callback to process the template + saved_callback = kw_copy.get('callback') + kw_copy['callback'] = lambda ctx, param, value: self._process_template(ctx, param, value, saved_callback) + + return click.option(*self.args, **kw_copy) + + def _process_template(self, ctx, param, value, saved_callback): + """Process the template file and update context defaults.""" + if not value: + return saved_callback(ctx, param, value) if saved_callback else value + + ctx.default_map = ctx.default_map or {} -# Create an enhanced config file option -TEMPLATE_FILE = OverridableOption( + # Get template vars from context if they were set by TEMPLATE_VARS option + template_vars = getattr(ctx, '_template_vars', None) + + # Determine if we're in non-interactive mode + non_interactive = False + for ctx_param in ctx.params.values(): + if isinstance(ctx_param, bool) and 'non_interactive' in str(ctx_param): + non_interactive = ctx_param + break + + try: + # Load and process the template + config_data = load_and_process_template(value, interactive=not non_interactive, template_vars=template_vars) + + # Update the default map with template values + # This will set defaults for options that haven't been explicitly provided + for key, template_value in config_data.items(): + if key not in ctx.default_map: + ctx.default_map[key] = template_value + + except Exception as e: + echo.echo_critical(f'Error processing template: {e}') + + return saved_callback(ctx, param, value) if saved_callback else value + + +# Updated TEMPLATE_FILE option +TEMPLATE_FILE = TemplateOption( '--template', type=click.STRING, help='Load computer setup from configuration file in YAML format (local path or URL). ' 'Supports Jinja2 templates with interactive prompting.', ) + +import json + + +# Helper function to set template vars in context (updated) +def set_template_vars_in_context(ctx, param, value): + """Set template variables in the context for the template option to use.""" + if value: + try: + template_var_dict = json.loads(value) + # Store template vars in context for the template option to use + ctx._template_vars = template_var_dict + except json.JSONDecodeError as e: + raise click.BadParameter(f'Invalid JSON in template-vars: {e}') + return value + + +# Updated TEMPLATE_VARS option with callback +TEMPLATE_VARS = OverridableOption( + '--template-vars', + type=click.STRING, + is_eager=True, # Process before template option + callback=set_template_vars_in_context, + expose_value=False, # Don't pass to command function + help='JSON string containing template variable values for non-interactive mode. ' + 'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'', +) From 65b378396338f99c7d030404f702586ca33406c4 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Wed, 18 Jun 2025 08:39:22 +0200 Subject: [PATCH 5/6] now it works --- src/aiida/cmdline/commands/cmd_computer.py | 120 +-------------------- src/aiida/cmdline/params/options/main.py | 108 ++++++++----------- 2 files changed, 47 insertions(+), 181 deletions(-) diff --git a/src/aiida/cmdline/commands/cmd_computer.py b/src/aiida/cmdline/commands/cmd_computer.py index 4fca0d503f..8ea2ccb492 100644 --- a/src/aiida/cmdline/commands/cmd_computer.py +++ b/src/aiida/cmdline/commands/cmd_computer.py @@ -285,8 +285,6 @@ def set_template_vars_in_context(ctx, param, value): return value - -# Modified computer setup command @verdi_computer.command('setup') @options_computer.LABEL() @options_computer.HOSTNAME() @@ -301,12 +299,12 @@ def set_template_vars_in_context(ctx, param, value): @options_computer.USE_DOUBLE_QUOTES() @options_computer.PREPEND_TEXT() @options_computer.APPEND_TEXT() -@options.NON_INTERACTIVE() # Uncomment this line -@options.TEMPLATE_FILE() -@options.TEMPLATE_VARS() +@options.NON_INTERACTIVE() +@options.TEMPLATE_VARS() # This should come before TEMPLATE_FILE +@options.TEMPLATE_FILE() # This will process the template and set defaults @click.pass_context @with_dbenv() -def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): +def computer_setup(ctx, non_interactive, **kwargs): """Create a new computer.""" from aiida.orm.utils.builders.computer import ComputerBuilder @@ -315,31 +313,6 @@ def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): print(f"Debug: kwargs keys = {list(kwargs.keys())}") print(f"Debug: ctx.default_map = {ctx.default_map}") - # Handle template variables - template_var_dict = None - if template_vars: - try: - template_var_dict = json.loads(template_vars) - except json.JSONDecodeError as e: - echo.echo_critical(f'Invalid JSON in template-vars: {e}') - - # Process template if provided - if template: - try: - # Load and process the template - config_data = load_and_process_template( - template, interactive=not non_interactive, template_vars=template_var_dict - ) - - # Update kwargs with config file data - # Only update if the value wasn't explicitly provided on command line - for key, value in config_data.items(): - if key not in kwargs or kwargs[key] is None: - kwargs[key] = value - - except Exception as e: - echo.echo_critical(f'Error processing template: {e}') - # Check for existing computer if kwargs.get('label') and kwargs['label'] in get_computer_names(): echo.echo_critical( @@ -372,91 +345,6 @@ def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): profile = ctx.obj['profile'] echo.echo_report(f' verdi -p {profile.name} computer configure {computer.transport_type} {computer.label}') -# Modified computer setup command -# @verdi_computer.command('setup') -# @options_computer.LABEL() -# @options_computer.HOSTNAME() -# @options_computer.DESCRIPTION() -# @options_computer.TRANSPORT() -# @options_computer.SCHEDULER() -# @options_computer.SHEBANG() -# @options_computer.WORKDIR() -# @options_computer.MPI_RUN_COMMAND() -# @options_computer.MPI_PROCS_PER_MACHINE() -# @options_computer.DEFAULT_MEMORY_PER_MACHINE() -# @options_computer.USE_DOUBLE_QUOTES() -# @options_computer.PREPEND_TEXT() -# @options_computer.APPEND_TEXT() -# # @options.NON_INTERACTIVE() -# # @options.CONFIG_FILE() # Keep the original config option for backward compatibility -# @options.TEMPLATE_FILE() # Add our new template option -# @options.TEMPLATE_VARS() -# @click.pass_context -# @with_dbenv() -# def computer_setup(ctx, non_interactive, template, template_vars, **kwargs): -# """Create a new computer.""" -# from aiida.orm.utils.builders.computer import ComputerBuilder - -# print('HELLO') - -# # Handle template variables -# template_var_dict = None -# if template_vars: -# try: -# template_var_dict = json.loads(template_vars) -# except json.JSONDecodeError as e: -# echo.echo_critical(f'Invalid JSON in template-vars: {e}') - -# # Process template if provided -# if template: -# try: -# # Load and process the template -# config_data = load_and_process_template( -# template, interactive=not non_interactive, template_vars=template_var_dict -# ) - -# # Update kwargs with config file data -# # Only update if the value wasn't explicitly provided on command line -# for key, value in config_data.items(): -# if key not in kwargs or kwargs[key] is None: -# kwargs[key] = value - -# except Exception as e: -# echo.echo_critical(f'Error processing template: {e}') - -# # Check for existing computer -# if kwargs.get('label') and kwargs['label'] in get_computer_names(): -# echo.echo_critical( -# 'A computer called {c} already exists. ' -# 'Use "verdi computer duplicate {c}" to set up a new ' -# 'computer starting from the settings of {c}.'.format(c=kwargs['label']) -# ) - -# # Convert entry points to their names -# if kwargs.get('transport'): -# kwargs['transport'] = kwargs['transport'].name -# if kwargs.get('scheduler'): -# kwargs['scheduler'] = kwargs['scheduler'].name - -# computer_builder = ComputerBuilder(**kwargs) -# try: -# computer = computer_builder.new() -# except (ComputerBuilder.ComputerValidationError, ValidationError) as e: -# echo.echo_critical(f'{type(e).__name__}: {e}') - -# try: -# computer.store() -# except ValidationError as err: -# echo.echo_critical(f'unable to store the computer: {err}. Exiting...') -# else: -# echo.echo_success(f'Computer<{computer.pk}> {computer.label} created') - -# echo.echo_report('Note: before the computer can be used, it has to be configured with the command:') - -# profile = ctx.obj['profile'] -# echo.echo_report(f' verdi -p {profile.name} computer configure {computer.transport_type} {computer.label}') - - @verdi_computer.command('duplicate') @arguments.COMPUTER(callback=set_computer_builder) @options_computer.LABEL(contextual_default=partial(get_parameter_default, 'label')) diff --git a/src/aiida/cmdline/params/options/main.py b/src/aiida/cmdline/params/options/main.py index 7e58174262..27e2696284 100644 --- a/src/aiida/cmdline/params/options/main.py +++ b/src/aiida/cmdline/params/options/main.py @@ -8,11 +8,14 @@ ########################################################################### """Module with pre-defined reusable commandline options that can be used as `click` decorators.""" +import json import pathlib import click from aiida.brokers.rabbitmq.defaults import BROKER_DEFAULTS +from aiida.cmdline.utils import echo +from aiida.cmdline.utils.template_config import load_and_process_template from aiida.common.log import LOG_LEVELS, configure_logging from aiida.manage.external.postgres import DEFAULT_DBINFO @@ -913,80 +916,46 @@ def set_log_level(ctx, _param, value): help='End date for node mtime range selection for node collection dumping.', ) +import json import click - from aiida.cmdline.utils import echo from aiida.cmdline.utils.template_config import load_and_process_template - from .overridable import OverridableOption +# Template processing callback +def process_template_callback(ctx, param, value): + """Process template file and update context defaults.""" + if not value: + return value -class TemplateOption(OverridableOption): - """Option that processes Jinja2 templates and updates Click context defaults.""" - - def __init__(self, *args, **kwargs): - """Initialize the template option.""" - kwargs.setdefault('is_eager', True) # Process template before other options - kwargs.setdefault('expose_value', False) # Don't pass template to the command function - super().__init__(*args, **kwargs) - - def __call__(self, **kwargs): - """Create the option with updated kwargs.""" - kw_copy = self.kwargs.copy() - kw_copy.update(kwargs) - - # Set the callback to process the template - saved_callback = kw_copy.get('callback') - kw_copy['callback'] = lambda ctx, param, value: self._process_template(ctx, param, value, saved_callback) - - return click.option(*self.args, **kw_copy) - - def _process_template(self, ctx, param, value, saved_callback): - """Process the template file and update context defaults.""" - if not value: - return saved_callback(ctx, param, value) if saved_callback else value - - ctx.default_map = ctx.default_map or {} - - # Get template vars from context if they were set by TEMPLATE_VARS option - template_vars = getattr(ctx, '_template_vars', None) - - # Determine if we're in non-interactive mode - non_interactive = False - for ctx_param in ctx.params.values(): - if isinstance(ctx_param, bool) and 'non_interactive' in str(ctx_param): - non_interactive = ctx_param - break - - try: - # Load and process the template - config_data = load_and_process_template(value, interactive=not non_interactive, template_vars=template_vars) + ctx.default_map = ctx.default_map or {} - # Update the default map with template values - # This will set defaults for options that haven't been explicitly provided - for key, template_value in config_data.items(): - if key not in ctx.default_map: - ctx.default_map[key] = template_value + # Get template vars from context if they were set by TEMPLATE_VARS option + template_vars = getattr(ctx, '_template_vars', None) - except Exception as e: - echo.echo_critical(f'Error processing template: {e}') + # Check if we're in non-interactive mode + non_interactive = ctx.params.get('non_interactive', False) - return saved_callback(ctx, param, value) if saved_callback else value + try: + # Load and process the template + config_data = load_and_process_template( + value, + interactive=not non_interactive, + template_vars=template_vars + ) + # Update the default map with template values + for key, template_value in config_data.items(): + if key not in ctx.default_map: + ctx.default_map[key] = template_value -# Updated TEMPLATE_FILE option -TEMPLATE_FILE = TemplateOption( - '--template', - type=click.STRING, - help='Load computer setup from configuration file in YAML format (local path or URL). ' - 'Supports Jinja2 templates with interactive prompting.', -) - -import json + except Exception as e: + echo.echo_critical(f'Error processing template: {e}') + return value -# Helper function to set template vars in context (updated) -def set_template_vars_in_context(ctx, param, value): +# Template vars callback +def set_template_vars_callback(ctx, param, value): """Set template variables in the context for the template option to use.""" if value: try: @@ -997,14 +966,23 @@ def set_template_vars_in_context(ctx, param, value): raise click.BadParameter(f'Invalid JSON in template-vars: {e}') return value - -# Updated TEMPLATE_VARS option with callback +# Template options using simple OverridableOption with callbacks TEMPLATE_VARS = OverridableOption( '--template-vars', type=click.STRING, is_eager=True, # Process before template option - callback=set_template_vars_in_context, + callback=set_template_vars_callback, expose_value=False, # Don't pass to command function help='JSON string containing template variable values for non-interactive mode. ' - 'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'', + 'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'', ) + +TEMPLATE_FILE = OverridableOption( + '--template', + type=click.STRING, + is_eager=True, # Process template before other options + callback=process_template_callback, + expose_value=False, # Don't pass template to the command function + help='Load computer setup from configuration file in YAML format (local path or URL). ' + 'Supports Jinja2 templates with interactive prompting.', +) \ No newline at end of file From 03f83fca0766aa9199e23604570ac85de099edb6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 18 Jun 2025 06:39:47 +0000 Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/aiida/cmdline/commands/cmd_computer.py | 8 ++++---- src/aiida/cmdline/params/options/main.py | 19 +++++++++---------- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/src/aiida/cmdline/commands/cmd_computer.py b/src/aiida/cmdline/commands/cmd_computer.py index 8ea2ccb492..7606daa10a 100644 --- a/src/aiida/cmdline/commands/cmd_computer.py +++ b/src/aiida/cmdline/commands/cmd_computer.py @@ -23,7 +23,6 @@ from aiida.cmdline.utils import echo, echo_tabulate from aiida.cmdline.utils.common import validate_output_filename from aiida.cmdline.utils.decorators import with_dbenv -from aiida.cmdline.utils.template_config import load_and_process_template from aiida.common.exceptions import EntryPointError, ValidationError from aiida.plugins.entry_point import get_entry_point_names @@ -309,9 +308,9 @@ def computer_setup(ctx, non_interactive, **kwargs): from aiida.orm.utils.builders.computer import ComputerBuilder # Debug output - print(f"Debug: non_interactive = {non_interactive}") - print(f"Debug: kwargs keys = {list(kwargs.keys())}") - print(f"Debug: ctx.default_map = {ctx.default_map}") + print(f'Debug: non_interactive = {non_interactive}') + print(f'Debug: kwargs keys = {list(kwargs.keys())}') + print(f'Debug: ctx.default_map = {ctx.default_map}') # Check for existing computer if kwargs.get('label') and kwargs['label'] in get_computer_names(): @@ -345,6 +344,7 @@ def computer_setup(ctx, non_interactive, **kwargs): profile = ctx.obj['profile'] echo.echo_report(f' verdi -p {profile.name} computer configure {computer.transport_type} {computer.label}') + @verdi_computer.command('duplicate') @arguments.COMPUTER(callback=set_computer_builder) @options_computer.LABEL(contextual_default=partial(get_parameter_default, 'label')) diff --git a/src/aiida/cmdline/params/options/main.py b/src/aiida/cmdline/params/options/main.py index 27e2696284..20daf9d3af 100644 --- a/src/aiida/cmdline/params/options/main.py +++ b/src/aiida/cmdline/params/options/main.py @@ -916,12 +916,13 @@ def set_log_level(ctx, _param, value): help='End date for node mtime range selection for node collection dumping.', ) -import json import click + from aiida.cmdline.utils import echo -from aiida.cmdline.utils.template_config import load_and_process_template + from .overridable import OverridableOption + # Template processing callback def process_template_callback(ctx, param, value): """Process template file and update context defaults.""" @@ -938,11 +939,7 @@ def process_template_callback(ctx, param, value): try: # Load and process the template - config_data = load_and_process_template( - value, - interactive=not non_interactive, - template_vars=template_vars - ) + config_data = load_and_process_template(value, interactive=not non_interactive, template_vars=template_vars) # Update the default map with template values for key, template_value in config_data.items(): @@ -954,6 +951,7 @@ def process_template_callback(ctx, param, value): return value + # Template vars callback def set_template_vars_callback(ctx, param, value): """Set template variables in the context for the template option to use.""" @@ -966,6 +964,7 @@ def set_template_vars_callback(ctx, param, value): raise click.BadParameter(f'Invalid JSON in template-vars: {e}') return value + # Template options using simple OverridableOption with callbacks TEMPLATE_VARS = OverridableOption( '--template-vars', @@ -974,7 +973,7 @@ def set_template_vars_callback(ctx, param, value): callback=set_template_vars_callback, expose_value=False, # Don't pass to command function help='JSON string containing template variable values for non-interactive mode. ' - 'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'', + 'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'', ) TEMPLATE_FILE = OverridableOption( @@ -984,5 +983,5 @@ def set_template_vars_callback(ctx, param, value): callback=process_template_callback, expose_value=False, # Don't pass template to the command function help='Load computer setup from configuration file in YAML format (local path or URL). ' - 'Supports Jinja2 templates with interactive prompting.', -) \ No newline at end of file + 'Supports Jinja2 templates with interactive prompting.', +)