Skip to content

Commit 65b3783

Browse files
GeigerJ2GeigerJ2
committed
now it works
1 parent 25488bb commit 65b3783

File tree

2 files changed

+47
-181
lines changed

2 files changed

+47
-181
lines changed

‎src/aiida/cmdline/commands/cmd_computer.py‎

Lines changed: 4 additions & 116 deletions
Original file line numberDiff line numberDiff line change
@@ -285,8 +285,6 @@ def set_template_vars_in_context(ctx, param, value):
285285
return value
286286

287287

288-
289-
# Modified computer setup command
290288
@verdi_computer.command('setup')
291289
@options_computer.LABEL()
292290
@options_computer.HOSTNAME()
@@ -301,12 +299,12 @@ def set_template_vars_in_context(ctx, param, value):
301299
@options_computer.USE_DOUBLE_QUOTES()
302300
@options_computer.PREPEND_TEXT()
303301
@options_computer.APPEND_TEXT()
304-
@options.NON_INTERACTIVE() # Uncomment this line
305-
@options.TEMPLATE_FILE()
306-
@options.TEMPLATE_VARS()
302+
@options.NON_INTERACTIVE()
303+
@options.TEMPLATE_VARS() # This should come before TEMPLATE_FILE
304+
@options.TEMPLATE_FILE() # This will process the template and set defaults
307305
@click.pass_context
308306
@with_dbenv()
309-
def computer_setup(ctx, non_interactive, template, template_vars, **kwargs):
307+
def computer_setup(ctx, non_interactive, **kwargs):
310308
"""Create a new computer."""
311309
from aiida.orm.utils.builders.computer import ComputerBuilder
312310

@@ -315,31 +313,6 @@ def computer_setup(ctx, non_interactive, template, template_vars, **kwargs):
315313
print(f"Debug: kwargs keys = {list(kwargs.keys())}")
316314
print(f"Debug: ctx.default_map = {ctx.default_map}")
317315

318-
# Handle template variables
319-
template_var_dict = None
320-
if template_vars:
321-
try:
322-
template_var_dict = json.loads(template_vars)
323-
except json.JSONDecodeError as e:
324-
echo.echo_critical(f'Invalid JSON in template-vars: {e}')
325-
326-
# Process template if provided
327-
if template:
328-
try:
329-
# Load and process the template
330-
config_data = load_and_process_template(
331-
template, interactive=not non_interactive, template_vars=template_var_dict
332-
)
333-
334-
# Update kwargs with config file data
335-
# Only update if the value wasn't explicitly provided on command line
336-
for key, value in config_data.items():
337-
if key not in kwargs or kwargs[key] is None:
338-
kwargs[key] = value
339-
340-
except Exception as e:
341-
echo.echo_critical(f'Error processing template: {e}')
342-
343316
# Check for existing computer
344317
if kwargs.get('label') and kwargs['label'] in get_computer_names():
345318
echo.echo_critical(
@@ -372,91 +345,6 @@ def computer_setup(ctx, non_interactive, template, template_vars, **kwargs):
372345
profile = ctx.obj['profile']
373346
echo.echo_report(f' verdi -p {profile.name} computer configure {computer.transport_type} {computer.label}')
374347

375-
# Modified computer setup command
376-
# @verdi_computer.command('setup')
377-
# @options_computer.LABEL()
378-
# @options_computer.HOSTNAME()
379-
# @options_computer.DESCRIPTION()
380-
# @options_computer.TRANSPORT()
381-
# @options_computer.SCHEDULER()
382-
# @options_computer.SHEBANG()
383-
# @options_computer.WORKDIR()
384-
# @options_computer.MPI_RUN_COMMAND()
385-
# @options_computer.MPI_PROCS_PER_MACHINE()
386-
# @options_computer.DEFAULT_MEMORY_PER_MACHINE()
387-
# @options_computer.USE_DOUBLE_QUOTES()
388-
# @options_computer.PREPEND_TEXT()
389-
# @options_computer.APPEND_TEXT()
390-
# # @options.NON_INTERACTIVE()
391-
# # @options.CONFIG_FILE() # Keep the original config option for backward compatibility
392-
# @options.TEMPLATE_FILE() # Add our new template option
393-
# @options.TEMPLATE_VARS()
394-
# @click.pass_context
395-
# @with_dbenv()
396-
# def computer_setup(ctx, non_interactive, template, template_vars, **kwargs):
397-
# """Create a new computer."""
398-
# from aiida.orm.utils.builders.computer import ComputerBuilder
399-
400-
# print('HELLO')
401-
402-
# # Handle template variables
403-
# template_var_dict = None
404-
# if template_vars:
405-
# try:
406-
# template_var_dict = json.loads(template_vars)
407-
# except json.JSONDecodeError as e:
408-
# echo.echo_critical(f'Invalid JSON in template-vars: {e}')
409-
410-
# # Process template if provided
411-
# if template:
412-
# try:
413-
# # Load and process the template
414-
# config_data = load_and_process_template(
415-
# template, interactive=not non_interactive, template_vars=template_var_dict
416-
# )
417-
418-
# # Update kwargs with config file data
419-
# # Only update if the value wasn't explicitly provided on command line
420-
# for key, value in config_data.items():
421-
# if key not in kwargs or kwargs[key] is None:
422-
# kwargs[key] = value
423-
424-
# except Exception as e:
425-
# echo.echo_critical(f'Error processing template: {e}')
426-
427-
# # Check for existing computer
428-
# if kwargs.get('label') and kwargs['label'] in get_computer_names():
429-
# echo.echo_critical(
430-
# 'A computer called {c} already exists. '
431-
# 'Use "verdi computer duplicate {c}" to set up a new '
432-
# 'computer starting from the settings of {c}.'.format(c=kwargs['label'])
433-
# )
434-
435-
# # Convert entry points to their names
436-
# if kwargs.get('transport'):
437-
# kwargs['transport'] = kwargs['transport'].name
438-
# if kwargs.get('scheduler'):
439-
# kwargs['scheduler'] = kwargs['scheduler'].name
440-
441-
# computer_builder = ComputerBuilder(**kwargs)
442-
# try:
443-
# computer = computer_builder.new()
444-
# except (ComputerBuilder.ComputerValidationError, ValidationError) as e:
445-
# echo.echo_critical(f'{type(e).__name__}: {e}')
446-
447-
# try:
448-
# computer.store()
449-
# except ValidationError as err:
450-
# echo.echo_critical(f'unable to store the computer: {err}. Exiting...')
451-
# else:
452-
# echo.echo_success(f'Computer<{computer.pk}> {computer.label} created')
453-
454-
# echo.echo_report('Note: before the computer can be used, it has to be configured with the command:')
455-
456-
# profile = ctx.obj['profile']
457-
# echo.echo_report(f' verdi -p {profile.name} computer configure {computer.transport_type} {computer.label}')
458-
459-
460348
@verdi_computer.command('duplicate')
461349
@arguments.COMPUTER(callback=set_computer_builder)
462350
@options_computer.LABEL(contextual_default=partial(get_parameter_default, 'label'))

‎src/aiida/cmdline/params/options/main.py‎

Lines changed: 43 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,14 @@
88
###########################################################################
99
"""Module with pre-defined reusable commandline options that can be used as `click` decorators."""
1010

11+
import json
1112
import pathlib
1213

1314
import click
1415

1516
from aiida.brokers.rabbitmq.defaults import BROKER_DEFAULTS
17+
from aiida.cmdline.utils import echo
18+
from aiida.cmdline.utils.template_config import load_and_process_template
1619
from aiida.common.log import LOG_LEVELS, configure_logging
1720
from aiida.manage.external.postgres import DEFAULT_DBINFO
1821

@@ -913,80 +916,46 @@ def set_log_level(ctx, _param, value):
913916
help='End date for node mtime range selection for node collection dumping.',
914917
)
915918

919+
import json
916920
import click
917-
918921
from aiida.cmdline.utils import echo
919922
from aiida.cmdline.utils.template_config import load_and_process_template
920-
921923
from .overridable import OverridableOption
922924

925+
# Template processing callback
926+
def process_template_callback(ctx, param, value):
927+
"""Process template file and update context defaults."""
928+
if not value:
929+
return value
923930

924-
class TemplateOption(OverridableOption):
925-
"""Option that processes Jinja2 templates and updates Click context defaults."""
926-
927-
def __init__(self, *args, **kwargs):
928-
"""Initialize the template option."""
929-
kwargs.setdefault('is_eager', True) # Process template before other options
930-
kwargs.setdefault('expose_value', False) # Don't pass template to the command function
931-
super().__init__(*args, **kwargs)
932-
933-
def __call__(self, **kwargs):
934-
"""Create the option with updated kwargs."""
935-
kw_copy = self.kwargs.copy()
936-
kw_copy.update(kwargs)
937-
938-
# Set the callback to process the template
939-
saved_callback = kw_copy.get('callback')
940-
kw_copy['callback'] = lambda ctx, param, value: self._process_template(ctx, param, value, saved_callback)
941-
942-
return click.option(*self.args, **kw_copy)
943-
944-
def _process_template(self, ctx, param, value, saved_callback):
945-
"""Process the template file and update context defaults."""
946-
if not value:
947-
return saved_callback(ctx, param, value) if saved_callback else value
948-
949-
ctx.default_map = ctx.default_map or {}
950-
951-
# Get template vars from context if they were set by TEMPLATE_VARS option
952-
template_vars = getattr(ctx, '_template_vars', None)
953-
954-
# Determine if we're in non-interactive mode
955-
non_interactive = False
956-
for ctx_param in ctx.params.values():
957-
if isinstance(ctx_param, bool) and 'non_interactive' in str(ctx_param):
958-
non_interactive = ctx_param
959-
break
960-
961-
try:
962-
# Load and process the template
963-
config_data = load_and_process_template(value, interactive=not non_interactive, template_vars=template_vars)
931+
ctx.default_map = ctx.default_map or {}
964932

965-
# Update the default map with template values
966-
# This will set defaults for options that haven't been explicitly provided
967-
for key, template_value in config_data.items():
968-
if key not in ctx.default_map:
969-
ctx.default_map[key] = template_value
933+
# Get template vars from context if they were set by TEMPLATE_VARS option
934+
template_vars = getattr(ctx, '_template_vars', None)
970935

971-
except Exception as e:
972-
echo.echo_critical(f'Error processing template: {e}')
936+
# Check if we're in non-interactive mode
937+
non_interactive = ctx.params.get('non_interactive', False)
973938

974-
return saved_callback(ctx, param, value) if saved_callback else value
939+
try:
940+
# Load and process the template
941+
config_data = load_and_process_template(
942+
value,
943+
interactive=not non_interactive,
944+
template_vars=template_vars
945+
)
975946

947+
# Update the default map with template values
948+
for key, template_value in config_data.items():
949+
if key not in ctx.default_map:
950+
ctx.default_map[key] = template_value
976951

977-
# Updated TEMPLATE_FILE option
978-
TEMPLATE_FILE = TemplateOption(
979-
'--template',
980-
type=click.STRING,
981-
help='Load computer setup from configuration file in YAML format (local path or URL). '
982-
'Supports Jinja2 templates with interactive prompting.',
983-
)
984-
985-
import json
952+
except Exception as e:
953+
echo.echo_critical(f'Error processing template: {e}')
986954

955+
return value
987956

988-
# Helper function to set template vars in context (updated)
989-
def set_template_vars_in_context(ctx, param, value):
957+
# Template vars callback
958+
def set_template_vars_callback(ctx, param, value):
990959
"""Set template variables in the context for the template option to use."""
991960
if value:
992961
try:
@@ -997,14 +966,23 @@ def set_template_vars_in_context(ctx, param, value):
997966
raise click.BadParameter(f'Invalid JSON in template-vars: {e}')
998967
return value
999968

1000-
1001-
# Updated TEMPLATE_VARS option with callback
969+
# Template options using simple OverridableOption with callbacks
1002970
TEMPLATE_VARS = OverridableOption(
1003971
'--template-vars',
1004972
type=click.STRING,
1005973
is_eager=True, # Process before template option
1006-
callback=set_template_vars_in_context,
974+
callback=set_template_vars_callback,
1007975
expose_value=False, # Don't pass to command function
1008976
help='JSON string containing template variable values for non-interactive mode. '
1009-
'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'',
977+
'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'',
1010978
)
979+
980+
TEMPLATE_FILE = OverridableOption(
981+
'--template',
982+
type=click.STRING,
983+
is_eager=True, # Process template before other options
984+
callback=process_template_callback,
985+
expose_value=False, # Don't pass template to the command function
986+
help='Load computer setup from configuration file in YAML format (local path or URL). '
987+
'Supports Jinja2 templates with interactive prompting.',
988+
)

0 commit comments

Comments
 (0)