PR comments and type changes

smartsim/_core/arguments/shell.py

@@ -32,17 +32,13 @@
 from smartsim.log import get_logger
 from smartsim.settings.arguments.launchArguments import LaunchArguments
 
-if t.TYPE_CHECKING:
-    from smartsim.settings.arguments import LaunchArguments
-
 logger = get_logger(__name__)
 
 
 class ShellLaunchArguments(LaunchArguments):
     @abstractmethod
     def format_env_vars(
-        self, env_vars: t.Dict[str, t.Optional[str]]
-    ) -> t.Union[t.List[str], None]: ...
-
+        self, env_vars: t.Mapping[str, str | None]
+    ) -> list[str] | None: ...
     @abstractmethod
-    def format_launch_args(self) -> t.Union[t.List[str], None]: ...
+    def format_launch_args(self) -> list[str] | None: ...

smartsim/_core/shell/shellLauncher.py

@@ -32,13 +32,12 @@
 import typing as t
 
 import psutil
-
+from smartsim._core.arguments.shell import ShellLaunchArguments
 from smartsim._core.dispatch import _EnvironMappingType, _FormatterType, dispatch
 from smartsim._core.utils import helpers
 from smartsim._core.utils.launcher import ExecutableProtocol, create_job_id
 from smartsim.error import errors
 from smartsim.log import get_logger
-from smartsim.settings.arguments.launchArguments import LaunchArguments
 from smartsim.status import JobStatus
 from smartsim.types import LaunchedJobID
 
@@ -100,11 +99,11 @@ def _get_status(self, id_: LaunchedJobID, /) -> JobStatus:
     def create(cls, _: Experiment) -> Self:
         return cls()
 
-    
+
 def make_shell_format_fn(
     run_command: str | None,
 ) -> _FormatterType[
-    LaunchArguments, tuple[str | os.PathLike[str], t.Sequence[str]]
+    ShellLaunchArguments, tuple[str | os.PathLike[str], t.Sequence[str]]
 ]:
     """A function that builds a function that formats a `LaunchArguments` as a
     shell executable sequence of strings for a given launching utility.
@@ -122,7 +121,7 @@ def make_shell_format_fn(
         as_srun_command = make_shell_format_fn("srun")
         fmt_cmd = as_srun_command(slurm_args, echo_hello_world, env)
         print(list(fmt_cmd))
-        # prints: "['srun', '--nodes=3', '--', 'echo', 'Hello World!']"
+        # prints: "['srun', --export ALL,foo=bar,baz=meep,spam=eggs '--nodes=3', '--', 'echo', 'Hello World!']"
 
     .. note::
         This function was/is a kind of slap-dash implementation, and is likely
@@ -136,7 +135,7 @@ def make_shell_format_fn(
     """
 
     def impl(
-        args: LaunchArguments,
+        args: ShellLaunchArguments,
         exe: ExecutableProtocol,
         path: str | os.PathLike[str],
         _env: _EnvironMappingType,

smartsim/settings/arguments/launch/alps.py

@@ -177,9 +177,7 @@ def set_quiet_launch(self, quiet: bool) -> None:
         else:
             self._launch_args.pop("quiet", None)
 
-    def format_env_vars(
-        self, env_vars: t.Optional[t.Dict[str, t.Optional[str]]]
-    ) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str] | None:
         """Format the environment variables for aprun
 
         :return: list of env vars

smartsim/settings/arguments/launch/local.py

@@ -49,7 +49,7 @@ def launcher_str(self) -> str:
         """
         return LauncherType.Local.value
 
-    def format_env_vars(self, env_vars: StringArgument) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str] | None:
         """Build bash compatible sequence of strings to specify an environment
 
         :param env_vars: An environment mapping

smartsim/settings/arguments/launch/lsf.py

@@ -74,9 +74,7 @@ def set_binding(self, binding: str) -> None:
         """
         self.set("bind", binding)
 
-    def format_env_vars(
-        self, env_vars: t.Dict[str, t.Optional[str]]
-    ) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str] | None:
         """Format environment variables. Each variable needs
         to be passed with ``--env``. If a variable is set to ``None``,
         its value is propagated from the current environment.

smartsim/settings/arguments/launch/mpi.py

@@ -172,9 +172,7 @@ def set_quiet_launch(self, quiet: bool) -> None:
         else:
             self._launch_args.pop("quiet", None)
 
-    def format_env_vars(
-        self, env_vars: t.Optional[t.Dict[str, t.Optional[str]]]
-    ) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str] | None:
         """Format the environment variables for mpirun
 
         :return: list of env vars

smartsim/settings/arguments/launch/pals.py

@@ -106,9 +106,7 @@ def set_hostlist(self, host_list: t.Union[str, t.List[str]]) -> None:
             raise TypeError("host_list argument must be list of strings")
         self.set("hosts", ",".join(host_list))
 
-    def format_env_vars(
-        self, env_vars: t.Optional[t.Dict[str, t.Optional[str]]]
-    ) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str] | None:
         """Format the environment variables for mpirun
 
         :return: list of env vars

smartsim/settings/arguments/launch/slurm.py

@@ -252,9 +252,7 @@ def format_launch_args(self) -> t.Union[t.List[str], None]:
                     formatted += ["=".join((prefix + key, str(value)))]
         return formatted
 
-    def format_env_vars(
-        self, env_vars: t.Dict[str, t.Optional[str]]
-    ) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str] | None:
         """Build bash compatible environment variable string for Slurm
 
         :returns: the formatted string of environment variables
@@ -291,7 +289,7 @@ def format_comma_sep_env_vars(
 
         return fmt_exported_env, compound_env
 
-    def _check_env_vars(self, env_vars: t.Dict[str, t.Optional[str]]) -> None:
+    def _check_env_vars(self, env_vars: t.Mapping[str, str | None]) -> None:
         """Warn a user trying to set a variable which is set in the environment
 
         Given Slurm's env var precedence, trying to export a variable which is already

smartsim/settings/launchSettings.py

@@ -29,7 +29,6 @@
 import copy
 import typing as t
 
-from smartsim._core.arguments.shell import ShellLaunchArguments
 from smartsim.log import get_logger
 
 from .._core.utils.helpers import fmt_dict
@@ -153,29 +152,6 @@ def update_env(self, env_vars: t.Dict[str, str | None]) -> None:
                 )
         self._env_vars.update(env_vars)
 
-    # def format_env_vars(self) -> t.Union[t.List[str], None]:
-    #     """Build bash compatible environment variable string for Slurm
-    #     :returns: the formatted string of environment variables
-    #     """
-    #     return self._arguments.format_env_vars(self._env_vars)
-
-    # def format_comma_sep_env_vars(self) -> t.Union[t.Tuple[str, t.List[str]], None]:
-    #     """Build environment variable string for Slurm
-    #     Slurm takes exports in comma separated lists
-    #     the list starts with all as to not disturb the rest of the environment
-    #     for more information on this, see the slurm documentation for srun
-    #     :returns: the formatted string of environment variables
-    #     """
-    #     return self._arguments.format_comma_sep_env_vars(self._env_vars)
-
-    # def format_launch_args(self) -> t.Union[t.List[str], None]:
-    #     """Return formatted launch arguments
-    #     For ``RunSettings``, the run arguments are passed
-    #     literally with no formatting.
-    #     :return: list run arguments for these settings
-    #     """
-    #     return self._arguments.format_launch_args()
-
     def __str__(self) -> str:  # pragma: no-cover
         string = f"\nLauncher: {self.launcher}{self.launch_args}"
         if self.env_vars: