Handle import errors from MACE

stfc · May 3, 2024 · bc176c9 · bc176c9
1 parent 1c36d1c
commit bc176c9
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/janus_core/cli/janus.py b/janus_core/cli/janus.py
@@ -8,7 +8,11 @@
 from janus_core.cli.geomopt import geomopt
 from janus_core.cli.md import md
 from janus_core.cli.singlepoint import singlepoint
-from janus_core.cli.train import train
+
+try:
+    from janus_core.cli.train import train
+except NotImplementedError:
+    pass
 
 app = Typer(name="janus", no_args_is_help=True)
 app.command()(singlepoint)

diff --git a/janus_core/helpers/train.py b/janus_core/helpers/train.py
@@ -3,7 +3,10 @@
 from pathlib import Path
 from typing import Optional
 
-from mace.cli.run_train import run as run_train
+try:
+    from mace.cli.run_train import run as run_train
+except ImportError as e:
+    raise NotImplementedError("Please update MACE to use this module.") from e
 from mace.tools import build_default_arg_parser as mace_parser
 import yaml