Use click choices in CLI (#493)

Author: ElliottKasoar

janus_core/cli/eos.py

@@ -4,6 +4,7 @@
 
 from typing import Annotated, get_args
 
+from click import Choice
 from typer import Context, Option, Typer
 from typer_config import use_config
 
@@ -21,6 +22,7 @@
     WriteKwargs,
 )
 from janus_core.cli.utils import yaml_converter_callback
+from janus_core.helpers.janus_types import EoSNames
 
 app = Typer()
 
@@ -35,7 +37,11 @@ def eos(
     max_volume: Annotated[float, Option(help="Maximum volume scale factor.")] = 1.05,
     n_volumes: Annotated[int, Option(help="Number of volumes.")] = 7,
     eos_type: Annotated[
-        str, Option(help="Type of fit for equation of state.")
+        str,
+        Option(
+            click_type=Choice(get_args(EoSNames)),
+            help="Type of fit for equation of state.",
+        ),
     ] = "birchmurnaghan",
     minimize: Annotated[
         bool, Option(help="Whether to minimize initial structure before calculations.")
@@ -140,7 +146,6 @@ def eos(
         set_read_kwargs_index,
         start_summary,
     )
-    from janus_core.helpers.janus_types import EoSNames
 
     # Check options from configuration file are all valid
     check_config(ctx)
@@ -149,9 +154,6 @@ def eos(
         [read_kwargs, calc_kwargs, minimize_kwargs, write_kwargs]
     )
 
-    if eos_type not in get_args(EoSNames):
-        raise ValueError(f"Fit type must be one of: {get_args(EoSNames)}")
-
     # Set initial config
     all_kwargs = {
         "read_kwargs": read_kwargs.copy(),

janus_core/cli/md.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 from typing import Annotated, get_args
 
+from click import Choice
 from typer import Context, Option, Typer
 from typer_config import use_config
 import yaml
@@ -27,6 +28,7 @@
     WriteKwargs,
 )
 from janus_core.cli.utils import parse_correlation_kwargs, yaml_converter_callback
+from janus_core.helpers.janus_types import Ensembles
 
 app = Typer()
 
@@ -58,7 +60,13 @@ def _update_restart_files(summary: Path, restart_files: list[Path]) -> None:
 def md(
     # numpydoc ignore=PR02
     ctx: Context,
-    ensemble: Annotated[str, Option(help="Name of thermodynamic ensemble.")],
+    ensemble: Annotated[
+        str,
+        Option(
+            click_type=Choice(get_args(Ensembles)),
+            help="Name of thermodynamic ensemble.",
+        ),
+    ],
     struct: StructPath,
     steps: Annotated[int, Option(help="Number of steps in simulation.")] = 0,
     timestep: Annotated[float, Option(help="Timestep for integrator, in fs.")] = 1.0,
@@ -380,7 +388,6 @@ def md(
         set_read_kwargs_index,
         start_summary,
     )
-    from janus_core.helpers.janus_types import Ensembles
 
     # Check options from configuration file are all valid
     check_config(ctx)
@@ -404,9 +411,6 @@ def md(
         ]
     )
 
-    if ensemble not in get_args(Ensembles):
-        raise ValueError(f"ensemble must be in {get_args(Ensembles)}")
-
     # Set initial config
     all_kwargs = {
         "read_kwargs": read_kwargs.copy(),

janus_core/cli/singlepoint.py

@@ -3,8 +3,9 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Annotated
+from typing import Annotated, get_args
 
+from click import Choice
 from typer import Context, Option, Typer
 from typer_config import use_config
 
@@ -21,6 +22,7 @@
     WriteKwargs,
 )
 from janus_core.cli.utils import yaml_converter_callback
+from janus_core.helpers.janus_types import Properties
 
 app = Typer()
 
@@ -37,6 +39,7 @@ def singlepoint(
     properties: Annotated[
         list[str] | None,
         Option(
+            click_type=Choice(get_args(Properties)),
             help=(
                 "Properties to calculate. If not specified, 'energy', 'forces' "
                 "and 'stress' will be returned."

janus_core/cli/types.py

@@ -4,10 +4,13 @@
 
 import ast
 from pathlib import Path
-from typing import TYPE_CHECKING, Annotated
+from typing import TYPE_CHECKING, Annotated, get_args
 
+from click import Choice
 from typer import Option
 
+from janus_core.helpers.janus_types import Architectures, Devices
+
 if TYPE_CHECKING:
     from janus_core.helpers.janus_types import ASEReadArgs
 
@@ -67,9 +70,19 @@ def __str__(self) -> str:
 StructPath = Annotated[Path, Option(help="Path of structure to simulate.")]
 
 Architecture = Annotated[
-    str | None, Option(help="MLIP architecture to use for calculations.")
+    str | None,
+    Option(
+        click_type=Choice(get_args(Architectures)),
+        help="MLIP architecture to use for calculations.",
+    ),
+]
+Device = Annotated[
+    str | None,
+    Option(
+        click_type=Choice(get_args(Devices)),
+        help="Device to run calculations on.",
+    ),
 ]
-Device = Annotated[str | None, Option(help="Device to run calculations on.")]
 ModelPath = Annotated[str | None, Option(help="Path to MLIP model.")]
 
 FilePrefix = Annotated[