7f8bd2da70
Co-authored-by: Claude Code <noreply@anthropic.com> GitOrigin-RevId: e1898de6e36887379f4bc44c4b03c08717615678
117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
"""Helper functions."""
|
|
|
|
import contextlib
|
|
import random
|
|
import re
|
|
import sys
|
|
from collections.abc import Generator
|
|
from typing import Any
|
|
|
|
import yaml
|
|
|
|
from buildscripts.resmokelib import config as _config
|
|
from buildscripts.resmokelib.utils import archival
|
|
|
|
__all__ = ["archival"]
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def open_or_use_stdout(filename: str) -> Generator[Any, None, None]:
|
|
"""Open the specified file for writing, or returns sys.stdout if filename is "-"."""
|
|
|
|
if filename == "-":
|
|
yield sys.stdout
|
|
return
|
|
|
|
line_buffered = 1
|
|
try:
|
|
fp = open(filename, "w", line_buffered)
|
|
except IOError:
|
|
print("Could not open file {}".format(filename), file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
try:
|
|
yield fp
|
|
finally:
|
|
fp.close()
|
|
|
|
|
|
def default_if_none(*values: Any) -> Any:
|
|
"""Return the first argument that is not 'None'."""
|
|
for value in values:
|
|
if value is not None:
|
|
return value
|
|
|
|
return None
|
|
|
|
|
|
def is_windows() -> bool:
|
|
"""Return True if Windows."""
|
|
return sys.platform.startswith("win32") or sys.platform.startswith("cygwin")
|
|
|
|
|
|
def is_string_list(lst: Any) -> bool:
|
|
"""Return true if 'lst' is a list of strings, and false otherwise."""
|
|
return isinstance(lst, list) and all(isinstance(x, str) for x in lst)
|
|
|
|
|
|
def load_yaml_file(filename: str) -> Any:
|
|
"""Attempt to read 'filename' as YAML."""
|
|
try:
|
|
with open(filename, "r", encoding="utf8") as fp:
|
|
return yaml.safe_load(fp)
|
|
except yaml.YAMLError as err:
|
|
raise ValueError("File '%s' contained invalid YAML: %s" % (filename, err))
|
|
|
|
|
|
def dump_yaml_file(value: Any, filename: str) -> None:
|
|
"""Attempt to write YAML object to filename."""
|
|
try:
|
|
with open(filename, "w") as fp:
|
|
return yaml.safe_dump(value, fp)
|
|
except yaml.YAMLError as err:
|
|
raise ValueError("Could not write YAML to file '%s': %s" % (filename, err))
|
|
|
|
|
|
def dump_yaml(value: Any) -> str:
|
|
"""Return 'value' formatted as YAML."""
|
|
# Use block (indented) style for formatting YAML.
|
|
return yaml.safe_dump(value, default_flow_style=False).rstrip()
|
|
|
|
|
|
def load_yaml(value: str) -> Any:
|
|
"""Attempt to parse 'value' as YAML."""
|
|
try:
|
|
return yaml.safe_load(value)
|
|
except yaml.YAMLError as err:
|
|
raise ValueError("Attempted to parse invalid YAML value '%s': %s" % (value, err))
|
|
|
|
|
|
def get_task_name_without_suffix(task_name: str, variant_name: str) -> str:
|
|
"""Return evergreen task name without suffix added to the generated task.
|
|
|
|
Remove evergreen variant name, numerical suffix and underscores between them from evergreen task name.
|
|
Example: "noPassthrough_0_enterprise-rhel-8-64-bit-dynamic-required" -> "noPassthrough"
|
|
"""
|
|
task_name = task_name if task_name else ""
|
|
return re.sub(rf"(_[0-9]+)?(_{variant_name})?$", "", task_name)
|
|
|
|
|
|
def pick_catalog_shard_node(config_shard: str | None, num_shards: int | None) -> int | None:
|
|
"""Get config_shard node index or None if no config_shard."""
|
|
if config_shard is None:
|
|
return None
|
|
|
|
if config_shard == "any":
|
|
# We check _config.NOOP_MONGO_D_S_PROCESSES because when running in antithesis
|
|
# the resmoke setup needs to be deterministic so the config shard cannot be random.
|
|
if num_shards is None or num_shards == 0 or _config.NOOP_MONGO_D_S_PROCESSES:
|
|
return 0
|
|
return random.randint(0, num_shards - 1)
|
|
|
|
config_shard_index = int(config_shard)
|
|
if config_shard_index < 0 or config_shard_index >= num_shards:
|
|
raise ValueError('Config shard value must be in range 0..num_shards-1 or "any"')
|
|
|
|
return config_shard_index
|