Payu general command wrapper for warning and exceptions formats

payu/cli.py

@@ -10,6 +10,7 @@
 import argparse
 import sysconfig
 import importlib
+import logging
 import os
 import pkgutil
 import shlex
@@ -24,23 +25,47 @@
 from payu.models import index as supported_models
 from payu.schedulers import index as scheduler_index, DEFAULT_SCHEDULER_CONFIG
 import payu.subcommands
+import payu.subcommands.args as arg_templates
 from payu.logger import setup_logger
 
+logger = logging.getLogger(__name__)
+
 # Default configuration
 DEFAULT_CONFIG = 'config.yaml'
 
-# Force warnings.warn() to omit the source code line in the message
-formatwarning_orig = warnings.formatwarning
-warnings.formatwarning = (
-    lambda message, category, filename, lineno, line=None: (
-        formatwarning_orig(message, category, filename, lineno, line='')
-    )
-)
+
+def _run_command(func, *args, stacktrace=False, log_level=None, **kwargs):
+    """Execute a payu command with error handling and logging.
+
+    Sets up logging, captures warnings through the logging system,
+    and catches exceptions to provide clean error messages.
+    """
+    setup_logger(log_level=log_level or logging.INFO)
+    # Capture warnings through the logging system
+    logging.captureWarnings(True)
+
+    if not stacktrace:
+        # Only display the warning message
+        warnings.formatwarning = (
+            lambda message, category, filename, lineno, line=None: str(message)
+        )
+
+    try:
+        func(*args, **kwargs)
+    except SystemExit:
+        # TODO: We want to remove any sys.exit(1) calls in the code
+        # and replace with exceptions! 
+        raise
+    except Exception as exc:
+        if stacktrace:
+            logger.exception(str(exc), exc_info=True)
+        else:
+            logger.error(str(exc))
+        sys.exit(1)
 
 
 def parse():
     """Parse the command line inputs and execute the subcommand."""
-    setup_logger()
     parser = generate_parser(is_interactive = True)
     # Display help if no arguments are provided
     if len(sys.argv) == 1:
@@ -50,7 +75,55 @@ def parse():
         parser = generate_parser()
     args = vars(parser.parse_args())
     run_cmd = args.pop('run_cmd')
-    run_cmd(**args)
+    stacktrace = args.pop('stacktrace')
+    log_level = args.pop('log_level')
+    _run_command(run_cmd, stacktrace=stacktrace, log_level=log_level, **args)
+
+
+# Add wrappers for runscript commands
+def parse_run():
+    _parse_runscript("run")
+
+
+def parse_collate():
+    _parse_runscript("collate")
+
+
+def parse_sync():
+    _parse_runscript("sync")
+
+
+def parse_profile():
+    _parse_runscript("profile")
+
+
+def _parse_runscript(cmd_name):
+
+    # Attempt to import the requested runscript command module
+    try:
+        cmd = importlib.import_module(f'payu.subcommands.{cmd_name}_cmd')
+    except ImportError:
+        print(f'payu: error: Unknown runscript command payu-{cmd_name}')
+        sys.exit(1)
+
+    # Construct the subcommand parser
+    parser = argparse.ArgumentParser(**cmd.parameters)
+
+    # Add global flags to each command
+    for arg in [arg_templates.stacktrace, arg_templates.log_level]:
+        parser.add_argument(*arg['flags'], **arg['parameters'])
+
+    for arg in cmd.arguments:
+        parser.add_argument(*arg['flags'], **arg['parameters'])
+
+    args = parser.parse_args()
+    stacktrace = args.stacktrace
+    log_level = args.log_level
+    # Remove them so they don't get passed to runscript
+    delattr(args, 'stacktrace')
+    delattr(args, 'log_level')
+
+    _run_command(cmd.runscript, args, stacktrace=stacktrace, log_level=log_level)
 
 
 def generate_parser(is_interactive=False):
@@ -75,6 +148,10 @@ def generate_parser(is_interactive=False):
         cmd_parser = subparsers.add_parser(cmd.title, **cmd.parameters)
         cmd_parser.set_defaults(run_cmd=cmd.runcmd)
 
+        # Add global flags to each subcommand
+        for arg in [arg_templates.stacktrace, arg_templates.log_level]:
+            cmd_parser.add_argument(*arg['flags'], **arg['parameters'])
+
         for arg in cmd.arguments:
             cmd_parser.add_argument(*arg['flags'], **arg['parameters'])
 

payu/logger.py

@@ -22,14 +22,14 @@ def format(self, record):
         formatter = logging.Formatter(log_color + self.FORMAT + self.RESET)
         return formatter.format(record)
 
-def setup_logger():
+def setup_logger(log_level=logging.INFO):
     """Configure the root logger"""
 
     # Color Formatter: Initialize Colorama for cross-platform compatibility and auto-reset
     init(autoreset=True)
 
     logger = logging.getLogger()
-    logger.setLevel(logging.INFO)
+    logger.setLevel(log_level)
 
     # Create a stream handler and set the custom formatter
     console_handler = logging.StreamHandler()

payu/subcommands/args.py

@@ -342,4 +342,24 @@
         'action': 'store',
         'help': 'Display information about a specific run number'
     }
+}
+
+stacktrace = {
+    'flags': ('--stacktrace',),
+    'parameters': {
+        'action': 'store_true',
+        'default': False,
+        'dest': 'stacktrace',
+        'help': 'Show full stack traces on errors and warnings'
+    }
+}
+
+log_level = {
+    'flags': ('--log-level',),
+    'parameters': {
+        'dest': 'log_level',
+        'choices': ['DEBUG', 'INFO', 'WARNING', 'ERROR'],
+        'default': None,
+        'help': 'Set logging verbosity level'
+    }
 }

payu/subcommands/collate_cmd.py

@@ -89,14 +89,7 @@ def runcmd(model_type, config_path, init_run, lab_path, dir_path):
     cli.submit_job('payu-collate', pbs_config, pbs_vars)
 
 
-def runscript():
-
-    parser = argparse.ArgumentParser()
-    for arg in arguments:
-        parser.add_argument(*arg['flags'], **arg['parameters'])
-
-    run_args = parser.parse_args()
-
+def runscript(run_args):
     pbs_vars = cli.set_env_vars(init_run=run_args.init_run,
                                 lab_path=run_args.lab_path,
                                 dir_path=run_args.dir_path)

payu/subcommands/profile_cmd.py

@@ -65,14 +65,7 @@ def runcmd(model_type, config_path, init_run, n_runs, lab_path):
     cli.submit_job('payu-profile', pbs_config, pbs_vars)
 
 
-def runscript():
-
-    parser = argparse.ArgumentParser()
-    for arg in arguments:
-        parser.add_argument(*arg['flags'], **arg['parameters'])
-
-    run_args = parser.parse_args()
-
+def runscript(run_args):
     pbs_vars = cli.set_env_vars(init_run=run_args.init_run,
                                 n_runs=run_args.n_runs)
     for var in pbs_vars:

payu/subcommands/run_cmd.py

@@ -189,13 +189,7 @@ def runcmd(model_type, config_path, init_run, n_runs, lab_path,
     )
 
 
-def runscript():
-    parser = argparse.ArgumentParser()
-    for arg in arguments:
-        parser.add_argument(*arg['flags'], **arg['parameters'])
-
-    run_args = parser.parse_args()
-
+def runscript(run_args):
     lab = Laboratory(run_args.model_type, run_args.config_path,
                      run_args.lab_path)
 

payu/subcommands/sync_cmd.py

@@ -61,13 +61,7 @@ def runcmd(model_type, config_path, lab_path, dir_path, sync_restarts,
     cli.submit_job('payu-sync', pbs_config, pbs_vars)
 
 
-def runscript():
-    parser = argparse.ArgumentParser()
-    for arg in arguments:
-        parser.add_argument(*arg['flags'], **arg['parameters'])
-
-    run_args = parser.parse_args()
-
+def runscript(run_args):
     pbs_vars = cli.set_env_vars(lab_path=run_args.lab_path,
                                 dir_path=run_args.dir_path,
                                 sync_restarts=run_args.sync_restarts,

pyproject.toml

@@ -54,13 +54,10 @@ mitgcm = ["mnctools>=0.2"]
 
 [project.scripts]
 payu = "payu.cli:parse"
-payu-run = "payu.subcommands.run_cmd:runscript"
-payu-collate = "payu.subcommands.collate_cmd:runscript"
-payu-profile = "payu.subcommands.profile_cmd:runscript"
-payu-sync = "payu.subcommands.sync_cmd:runscript"
-payu-branch = "payu.subcommands.branch_cmd:runscript"
-payu-clone = "payu.subcommands.clone_cmd:runscript"
-payu-checkout = "payu.subcommands.checkout_cmd:runscript"
+payu-run = "payu.cli:parse_run"
+payu-collate = "payu.cli:parse_collate"
+payu-profile = "payu.cli:parse_profile"
+payu-sync = "payu.cli:parse_sync"
 
 [build-system]
 build-backend = "setuptools.build_meta"