from __future__ import annotations
import os
import shutil
import subprocess
import warnings
from pathlib import Path
from monty.os import cd
from jobflow_remote.remote.host.base import BaseHost
[docs]
class LocalHost(BaseHost):
def __init__(self, timeout_execute: int = None) -> None:
self.timeout_execute = timeout_execute
def __eq__(self, other):
return isinstance(other, LocalHost)
[docs]
def execute(
self,
command: str | list[str],
workdir: str | Path | None = None,
timeout: int | None = None,
):
"""Execute the given command on the host.
Note that the command is executed with shell=True, so commands can
be exposed to command injection. Consider whether to escape part of
the input if it comes from external users.
Parameters
----------
command: str or list of str
Command to execute, as a str or list of str
Returns
-------
stdout : str
Standard output of the command
stderr : str
Standard error of the command
exit_code : int
Exit code of the command.
"""
if isinstance(command, (list, tuple)):
command = " ".join(command)
workdir = str(workdir) if workdir else Path.cwd()
timeout = timeout or self.timeout_execute
with cd(workdir):
proc = subprocess.run(
command, capture_output=True, shell=True, timeout=timeout, check=False
)
return proc.stdout.decode(), proc.stderr.decode(), proc.returncode
[docs]
def mkdir(
self, directory: str | Path, recursive: bool = True, exist_ok: bool = True
) -> bool:
try:
Path(directory).mkdir(parents=recursive, exist_ok=exist_ok)
except OSError:
return False
return True
[docs]
def write_text_file(self, filepath, content) -> None:
Path(filepath).write_text(content)
[docs]
def connect(self) -> None:
pass
[docs]
def close(self) -> bool:
return True
@property
def is_connected(self) -> bool:
return True
[docs]
def put(self, src, dst) -> None:
is_file_like = hasattr(src, "read") and callable(src.read)
src_base = getattr(src, "name", None) if is_file_like else os.path.basename(src)
if Path(dst).is_dir():
if src_base:
dst = Path(dst, src_base)
elif is_file_like:
raise ValueError(
"could not determine the file name and dst is a folder"
)
if is_file_like:
with open(dst, "wb") as f:
f.write(src.read())
else:
self.copy(src, dst)
[docs]
def get(self, src, dst) -> None:
is_file_like = hasattr(dst, "write") and callable(dst.write)
if is_file_like:
with open(src, "rb") as f:
dst.write(f.read())
else:
self.copy(src, dst)
[docs]
def copy(self, src, dst) -> None:
shutil.copy(src, dst)
[docs]
def listdir(self, path: str | Path) -> list[str]:
try:
return os.listdir(path)
except FileNotFoundError:
return []
[docs]
def remove(self, path: str | Path) -> None:
os.remove(path)
[docs]
def rmtree(self, path: str | Path, raise_on_error: bool = False) -> bool:
"""Recursively delete a directory tree on a local host.
It is intended to remove an entire directory tree, including all files
and subdirectories, on this local host.
Parameters
----------
path : str or Path
The path to the directory tree to be removed.
raise_on_error : bool
If set to `False` (default), errors will be ignored, and the method will
attempt to continue removing remaining files and directories.
Otherwise, any errors encountered during the removal process
will raise an exception.
Returns
-------
bool
True if the directory tree was successfully removed, False otherwise.
"""
removed = True
def onerror(func, dir_path, err):
nonlocal removed
removed = False
warnings.warn(f"Error while deleting folder {path}: {err[1]}", stacklevel=2)
shutil.rmtree(path, onerror=onerror if not raise_on_error else None)
return removed