def pytest_clean_artifacts(dir_name: str, preview: bool = False) -> None:
_LOG.warning("Cleaning pytest artifacts")
dbg.dassert_ne(dir_name, "")
dbg.dassert_dir_exists(dir_name)
if preview:
_LOG.warning("Preview only: nothing will be deleted")
# Show before cleaning.
file_names = pytest_show_artifacts(dir_name, tag="Before cleaning")
# Clean.
for f in file_names:
exists = os.path.exists(f)
_LOG.debug("%s -> exists=%s", f, exists)
if exists:
if not preview:
if os.path.isdir(f):
shutil.rmtree(f)
elif os.path.isfile(f):
os.remove(f)
else:
raise ValueError("Can't delete %s" % f)
else:
_LOG.debug("rm %s", f)
# Show after cleaning.
file_names = pytest_show_artifacts(dir_name, tag="After cleaning")
dbg.dassert_eq(len(file_names), 0)
def frame(
message: str,
char1: Optional[str] = None,
num_chars: Optional[int] = None,
char2: Optional[str] = None,
thickness: int = 1,
) -> str:
"""Print a frame around a message."""
# Fill in the default values.
if char1 is None:
# User didn't specify any char.
char1 = char2 = "#"
elif char1 is not None and char2 is None:
# User specified only one char.
char2 = char1
elif char1 is None and char2 is not None:
# User specified the second char, but not the first.
dbg.dfatal("Invalid char1='%s' char2='%s'" % (char1, char2))
else:
# User specified both chars. Nothing to do.
pass
num_chars = 80 if num_chars is None else num_chars
# Sanity check.
dbg.dassert_lte(1, thickness)
dbg.dassert_eq(len(char1), 1)
dbg.dassert_eq(len(char2), 1)
dbg.dassert_lte(1, num_chars)
# Build the return value.
ret = ((line(char1, num_chars) + "\n") * thickness + message + "\n" +
(line(char2, num_chars) + "\n") * thickness).rstrip("\n")
return ret
def kill_process(
get_pids: Callable[[], Tuple[List[int], str]],
timeout_in_secs: int = 5,
polltime_in_secs: float = 0.1,
) -> None:
"""Kill all the processes returned by the function `get_pids()`.
:param timeout_in_secs: how many seconds to wait at most before giving up
:param polltime_in_secs: how often to check for dead processes
"""
pids, txt = get_pids()
_LOG.info("Killing %d pids (%s)\n%s", len(pids), pids, "\n".join(txt))
if not pids:
return
for pid in pids:
try:
os.kill(pid, signal.SIGKILL)
except ProcessLookupError as e:
_LOG.warning(str(e))
#
_LOG.info("Waiting %d processes (%s) to die", len(pids), pids)
import tqdm
for _ in tqdm.tqdm(range(int(timeout_in_secs / polltime_in_secs))):
time.sleep(polltime_in_secs)
pids, _ = get_pids()
if not pids:
break
pids, txt = get_pids()
dbg.dassert_eq(len(pids), 0, "Processes are still alive:%s",
"\n".join(txt))
_LOG.info("Processes dead")
def get_repo_symbolic_name_from_dirname(git_dir: str) -> str:
"""Return the name of the repo in `git_dir`.
E.g., "alphamatic/amp", "ParticleDev/commodity_research"
"""
dbg.dassert_exists(git_dir)
cmd = "cd %s; (git remote -v | grep fetch)" % git_dir
# TODO(gp): Make it more robust, by checking both fetch and push.
# "origin [email protected]:alphamatic/amp (fetch)"
_, output = si.system_to_string(cmd)
data: List[str] = output.split()
_LOG.debug("data=%s", data)
dbg.dassert_eq(len(data), 3, "data='%s'", str(data))
# [email protected]:alphamatic/amp
repo_name = data[1]
m = re.match(r"^.*\.com:(.*)$", repo_name)
dbg.dassert(m, "Can't parse '%s'", repo_name)
repo_name = m.group(1) # type: ignore
_LOG.debug("repo_name=%s", repo_name)
# We expect something like "alphamatic/amp".
m = re.match(r"^\S+/\S+$", repo_name)
dbg.dassert(m, "repo_name='%s'", repo_name)
# origin [email protected]:ParticleDev/ORG_Particle.git (fetch)
suffix_to_remove = ".git"
if repo_name.endswith(suffix_to_remove):
repo_name = repo_name[:-len(suffix_to_remove)]
return repo_name
def get_first_line(output: str) -> str:
"""Return the first (and only) line from a string.
This is used when calling system_to_string() and expecting a single
line output.
"""
output = prnt.remove_empty_lines(output)
output_as_arr: List[str] = output.split("\n")
dbg.dassert_eq(len(output_as_arr), 1, "output='%s'", output)
output = output_as_arr[0]
output = output.rstrip().lstrip()
return output
def validate_datetime(timestamp: DATETIME_TYPE) -> pd.Timestamp:
"""
Assert that timestamp is in UTC, convert to pd.Timestamp.
:param timestamp: datetime object or pd.Timestamp
:return: tz-aware pd.Timestamp
"""
dbg.dassert_type_in(timestamp, [pd.Timestamp, datetime.datetime])
pd_timestamp = pd.Timestamp(timestamp)
dbg.dassert(pd_timestamp.tzinfo, "Timestamp should be tz-aware.")
dbg.dassert_eq(pd_timestamp.tzinfo.zone, "UTC", "Timezone should be UTC.")
return pd_timestamp
def get_client_root(super_module: bool) -> str:
"""Return the full path of the root of the Git client.
E.g., `/Users/saggese/src/.../amp`.
:param super_module: if True use the root of the Git _super_module,
if we are in a submodule. Otherwise use the Git _sub_module root
"""
if super_module and is_inside_submodule():
# https://stackoverflow.com/questions/957928
cmd = "git rev-parse --show-superproject-working-tree"
else:
cmd = "git rev-parse --show-toplevel"
_, out = si.system_to_string(cmd)
out = out.rstrip("\n")
dbg.dassert_eq(len(out.split("\n")), 1, msg="Invalid out='%s'" % out)
client_root: str = os.path.realpath(out)
return client_root
def get_amp_abs_path() -> str:
"""Return the absolute path of `amp` dir."""
repo_sym_name = get_repo_symbolic_name(super_module=False)
if repo_sym_name == "alphamatic/amp":
# If we are in the amp repo, then the git client root is the amp
# directory.
git_root = get_client_root(super_module=False)
amp_dir = git_root
else:
# If we are not in the amp repo, then look for the amp dir.
amp_dir = find_file_in_git_tree("amp", super_module=True)
git_root = get_client_root(super_module=True)
amp_dir = os.path.join(git_root, amp_dir)
amp_dir = os.path.abspath(amp_dir)
# Sanity check.
dbg.dassert_dir_exists(amp_dir)
if si.get_user_name() != "jenkins":
# Jenkins checks out amp repo in directories with different names,
# e.g., amp.dev.build_clean_env.run_slow_coverage_tests.
dbg.dassert_eq(os.path.basename(amp_dir), "amp")
return amp_dir
def _system(
cmd: str,
abort_on_error: bool,
suppress_error: Optional[Any],
suppress_output: bool,
blocking: bool,
wrapper: Optional[Any],
output_file: Optional[Any],
tee: bool,
dry_run: bool,
log_level: Union[int, str],
) -> Tuple[int, str]:
"""Execute a shell command.
:param cmd: string with command to execute
:param abort_on_error: whether we should assert in case of error or not
:param suppress_error: set of error codes to suppress
:param suppress_output: whether to print the output or not
- If "on_debug_level" then print the output if the log level is DEBUG
:param blocking: blocking system call or not
:param wrapper: another command to prepend the execution of cmd
:param output_file: redirect stdout and stderr to this file
:param tee: if True, tee stdout and stderr to output_file
:param dry_run: just print the final command but not execute it
:param log_level: print the command to execute at level "log_level".
- If "echo" then print the command line to screen as print and not
logging
:return: return code (int), output of the command (str)
"""
orig_cmd = cmd[:]
# Prepare the command line.
cmd = "(%s)" % cmd
dbg.dassert_imply(tee, output_file is not None)
if output_file is not None:
dir_name = os.path.dirname(output_file)
if not os.path.exists(dir_name):
_LOG.debug("Dir '%s' doesn't exist: creating", dir_name)
dbg.dassert(bool(dir_name), "dir_name='%s'", dir_name)
os.makedirs(dir_name)
if tee:
cmd += " 2>&1 | tee %s" % output_file
else:
cmd += " 2>&1 >%s" % output_file
else:
cmd += " 2>&1"
if wrapper:
cmd = wrapper + " && " + cmd
#
# TODO(gp): Add a check for the valid values.
# TODO(gp): Make it "ECHO".
if isinstance(log_level, str):
dbg.dassert_eq(log_level, "echo")
print("> %s" % orig_cmd)
_LOG.debug("> %s", cmd)
else:
_LOG.log(log_level, "> %s", cmd)
#
dbg.dassert_in(suppress_output, ("ON_DEBUG_LEVEL", True, False))
if suppress_output == "ON_DEBUG_LEVEL":
# print("eff_lev=%s" % eff_level)
# print("lev=%s" % logging.DEBUG)
_LOG.getEffectiveLevel()
# Suppress the output if the verbosity level is higher than DEBUG,
# otherwise print.
suppress_output = _LOG.getEffectiveLevel() > logging.DEBUG
#
output = ""
if dry_run:
_LOG.warning("Not executing cmd\n%s\nas per user request", cmd)
rc = 0
return rc, output
# Execute the command.
try:
stdout = subprocess.PIPE
stderr = subprocess.STDOUT
p = subprocess.Popen(cmd,
shell=True,
executable="/bin/bash",
stdout=stdout,
stderr=stderr)
output = ""
if blocking:
# Blocking call: get the output.
while True:
line = p.stdout.readline().decode("utf-8") # type: ignore
if not line:
break
if not suppress_output:
print((line.rstrip("\n")))
output += line
p.stdout.close() # type: ignore
rc = p.wait()
else:
# Not blocking.
# Wait until process terminates (without using p.wait()).
max_cnt = 20
cnt = 0
while p.poll() is None:
# Process hasn't exited yet, let's wait some time.
time.sleep(0.1)
cnt += 1
_LOG.debug("cnt=%s, rc=%s", cnt, p.returncode)
if cnt > max_cnt:
break
if cnt > max_cnt:
# Timeout: we assume it worked.
rc = 0
else:
rc = p.returncode
if suppress_error is not None:
dbg.dassert_isinstance(suppress_error, set)
if rc in suppress_error:
rc = 0
except OSError as e:
rc = -1
_LOG.error("error=%s", str(e))
_LOG.debug("rc=%s", rc)
if abort_on_error and rc != 0:
msg = ("\n" + prnt.frame("cmd='%s' failed with rc='%s'" % (cmd, rc)) +
"\nOutput of the failing command is:\n%s\n%s\n%s" %
(prnt.line(">"), output, prnt.line("<")))
_LOG.error("%s", msg)
raise RuntimeError("cmd='%s' failed with rc='%s'" % (cmd, rc))
# dbg.dassert_type_in(output, (str, ))
return rc, output