2023-10-03 14:47:14 +00:00
|
|
|
import logging
|
2024-01-09 16:34:43 +00:00
|
|
|
import os
|
|
|
|
import select
|
2023-10-03 14:47:14 +00:00
|
|
|
import shlex
|
2024-01-02 15:34:10 +00:00
|
|
|
import subprocess
|
|
|
|
import sys
|
2024-03-08 15:01:54 +00:00
|
|
|
import weakref
|
|
|
|
from datetime import datetime, timedelta
|
2024-01-03 13:25:34 +00:00
|
|
|
from enum import Enum
|
2023-10-09 12:01:34 +00:00
|
|
|
from pathlib import Path
|
2024-06-03 10:23:56 +00:00
|
|
|
from typing import IO, Any
|
2023-10-03 14:47:14 +00:00
|
|
|
|
2024-01-11 20:48:39 +00:00
|
|
|
from .custom_logger import get_caller
|
2024-01-10 16:58:39 +00:00
|
|
|
from .errors import ClanCmdError, CmdOut
|
2023-10-03 14:47:14 +00:00
|
|
|
|
2024-01-11 20:48:39 +00:00
|
|
|
glog = logging.getLogger(__name__)
|
2023-10-03 14:47:14 +00:00
|
|
|
|
|
|
|
|
2024-01-10 16:18:21 +00:00
|
|
|
class Log(Enum):
|
|
|
|
STDERR = 1
|
|
|
|
STDOUT = 2
|
|
|
|
BOTH = 3
|
2024-01-11 23:08:15 +00:00
|
|
|
NONE = 4
|
2024-01-10 16:18:21 +00:00
|
|
|
|
2024-01-11 23:13:21 +00:00
|
|
|
|
2024-06-03 10:23:56 +00:00
|
|
|
def handle_output(process: subprocess.Popen, log: Log) -> tuple[str, str]:
|
|
|
|
rlist = [process.stdout, process.stderr]
|
|
|
|
stdout_buf = b""
|
|
|
|
stderr_buf = b""
|
|
|
|
|
|
|
|
while len(rlist) != 0:
|
|
|
|
r, _, _ = select.select(rlist, [], [], 0.1)
|
|
|
|
if len(r) == 0: # timeout in select
|
|
|
|
if process.poll() is None:
|
|
|
|
continue
|
|
|
|
# Process has exited
|
|
|
|
break
|
|
|
|
|
|
|
|
def handle_fd(fd: IO[Any] | None) -> bytes:
|
|
|
|
if fd and fd in r:
|
|
|
|
read = os.read(fd.fileno(), 4096)
|
|
|
|
if len(read) != 0:
|
|
|
|
return read
|
|
|
|
rlist.remove(fd)
|
|
|
|
return b""
|
|
|
|
|
|
|
|
ret = handle_fd(process.stdout)
|
|
|
|
if ret and log in [Log.STDOUT, Log.BOTH]:
|
|
|
|
sys.stdout.buffer.write(ret)
|
|
|
|
sys.stdout.flush()
|
|
|
|
|
|
|
|
stdout_buf += ret
|
|
|
|
ret = handle_fd(process.stderr)
|
|
|
|
|
|
|
|
if ret and log in [Log.STDERR, Log.BOTH]:
|
|
|
|
sys.stderr.buffer.write(ret)
|
|
|
|
sys.stderr.flush()
|
|
|
|
stderr_buf += ret
|
|
|
|
return stdout_buf.decode("utf-8", "replace"), stderr_buf.decode("utf-8", "replace")
|
|
|
|
|
|
|
|
|
2024-03-08 15:01:54 +00:00
|
|
|
class TimeTable:
|
|
|
|
"""
|
|
|
|
This class is used to store the time taken by each command
|
|
|
|
and print it at the end of the program if env PERF=1 is set.
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self) -> None:
|
|
|
|
self.table: dict[str, timedelta] = {}
|
|
|
|
weakref.finalize(self, self.table_print)
|
|
|
|
|
|
|
|
def table_print(self) -> None:
|
|
|
|
if os.getenv("PERF") != "1":
|
|
|
|
return
|
|
|
|
print("======== CMD TIMETABLE ========")
|
|
|
|
|
|
|
|
# Sort the table by time in descending order
|
|
|
|
sorted_table = sorted(
|
|
|
|
self.table.items(), key=lambda item: item[1], reverse=True
|
|
|
|
)
|
|
|
|
|
|
|
|
for k, v in sorted_table:
|
|
|
|
# Check if timedelta is greater than 1 second
|
|
|
|
if v.total_seconds() > 1:
|
|
|
|
# Print in red
|
|
|
|
print(f"\033[91mTook {v}s\033[0m for command: '{k}'")
|
|
|
|
else:
|
|
|
|
# Print in default color
|
|
|
|
print(f"Took {v} for command: '{k}'")
|
|
|
|
|
|
|
|
def add(self, cmd: str, time: timedelta) -> None:
|
|
|
|
if cmd in self.table:
|
|
|
|
self.table[cmd] += time
|
|
|
|
else:
|
|
|
|
self.table[cmd] = time
|
|
|
|
|
|
|
|
|
|
|
|
TIME_TABLE = TimeTable()
|
|
|
|
|
|
|
|
|
2024-01-10 16:58:39 +00:00
|
|
|
def run(
|
|
|
|
cmd: list[str],
|
|
|
|
*,
|
2024-06-02 07:53:24 +00:00
|
|
|
input: bytes | None = None, # noqa: A002
|
2024-01-10 16:58:39 +00:00
|
|
|
env: dict[str, str] | None = None,
|
|
|
|
cwd: Path = Path.cwd(),
|
|
|
|
log: Log = Log.STDERR,
|
2024-01-10 17:46:54 +00:00
|
|
|
check: bool = True,
|
2024-01-10 18:29:16 +00:00
|
|
|
error_msg: str | None = None,
|
2024-01-10 16:58:39 +00:00
|
|
|
) -> CmdOut:
|
2024-06-02 07:53:24 +00:00
|
|
|
if input:
|
|
|
|
glog.debug(
|
2024-06-02 08:00:19 +00:00
|
|
|
f"""$: echo "{input.decode('utf-8', 'replace')}" | {shlex.join(cmd)} \nCaller: {get_caller()}"""
|
2024-06-02 07:53:24 +00:00
|
|
|
)
|
|
|
|
else:
|
|
|
|
glog.debug(f"$: {shlex.join(cmd)} \nCaller: {get_caller()}")
|
2024-03-06 10:12:20 +00:00
|
|
|
tstart = datetime.now()
|
|
|
|
|
2024-06-03 10:23:56 +00:00
|
|
|
# Start the subprocess
|
|
|
|
process = subprocess.Popen(
|
2024-01-10 16:58:39 +00:00
|
|
|
cmd,
|
2024-06-02 14:52:31 +00:00
|
|
|
cwd=str(cwd),
|
2024-06-03 10:23:56 +00:00
|
|
|
env=env,
|
|
|
|
stdout=subprocess.PIPE,
|
|
|
|
stderr=subprocess.PIPE,
|
2023-10-03 14:47:14 +00:00
|
|
|
)
|
2024-06-03 10:23:56 +00:00
|
|
|
stdout_buf, stderr_buf = handle_output(process, log)
|
2024-01-11 23:08:15 +00:00
|
|
|
|
2024-06-02 07:53:24 +00:00
|
|
|
if input:
|
2024-06-03 10:23:56 +00:00
|
|
|
process.communicate(input)
|
|
|
|
else:
|
|
|
|
process.wait()
|
2024-06-02 14:52:31 +00:00
|
|
|
tend = datetime.now()
|
2024-06-03 10:23:56 +00:00
|
|
|
|
2024-03-08 15:01:54 +00:00
|
|
|
global TIME_TABLE
|
|
|
|
TIME_TABLE.add(shlex.join(cmd), tend - tstart)
|
2024-01-11 23:16:02 +00:00
|
|
|
|
|
|
|
# Wait for the subprocess to finish
|
2024-01-10 16:58:39 +00:00
|
|
|
cmd_out = CmdOut(
|
2024-06-03 10:23:56 +00:00
|
|
|
stdout=stdout_buf,
|
|
|
|
stderr=stderr_buf,
|
2024-01-10 16:58:39 +00:00
|
|
|
cwd=cwd,
|
|
|
|
command=shlex.join(cmd),
|
2024-06-03 10:23:56 +00:00
|
|
|
returncode=process.returncode,
|
2024-01-10 18:29:16 +00:00
|
|
|
msg=error_msg,
|
2024-01-10 16:58:39 +00:00
|
|
|
)
|
2024-01-02 15:53:01 +00:00
|
|
|
|
2024-06-03 10:23:56 +00:00
|
|
|
if check and process.returncode != 0:
|
2024-01-10 16:58:39 +00:00
|
|
|
raise ClanCmdError(cmd_out)
|
|
|
|
|
|
|
|
return cmd_out
|
2024-05-28 09:13:55 +00:00
|
|
|
|
|
|
|
|
|
|
|
def run_no_stdout(
|
|
|
|
cmd: list[str],
|
|
|
|
*,
|
|
|
|
env: dict[str, str] | None = None,
|
|
|
|
cwd: Path = Path.cwd(),
|
|
|
|
log: Log = Log.STDERR,
|
|
|
|
check: bool = True,
|
|
|
|
error_msg: str | None = None,
|
|
|
|
) -> CmdOut:
|
|
|
|
"""
|
|
|
|
Like run, but automatically suppresses stdout, if not in DEBUG log level.
|
|
|
|
If in DEBUG log level the stdout of commands will be shown.
|
|
|
|
"""
|
|
|
|
if logging.getLogger(__name__.split(".")[0]).isEnabledFor(logging.DEBUG):
|
|
|
|
return run(cmd, env=env, log=log, check=check, error_msg=error_msg)
|
|
|
|
else:
|
|
|
|
log = Log.NONE
|
|
|
|
return run(cmd, env=env, log=log, check=check, error_msg=error_msg)
|