thevar1able/ClickHouse
1
0
Fork 0
mirror of https://github.com/ClickHouse/ClickHouse.git synced 2024-11-10 09:32:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Improve TeePopen typing and structure

Browse Source
This commit is contained in:
Mikhail f. Shiryaev 2022-12-14 11:53:17 +01:00
parent fdecb85dc4
commit 16b2d5dfd1
No known key found for this signature in database
GPG Key ID: 4B02ED204C7D93F4
1 changed files with 27 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
tests/ci/tee_popen.py
View File

@ -1,5 +1,6 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
from io import TextIOWrapper
from subprocess import Popen, PIPE, STDOUT from subprocess import Popen, PIPE, STDOUT
from threading import Thread from threading import Thread
from time import sleep from time import sleep
@ -14,15 +15,23 @@ import sys
# it finishes. stderr and stdout will be redirected both to specified file and # it finishes. stderr and stdout will be redirected both to specified file and
# stdout. # stdout.
class TeePopen: class TeePopen:
# pylint: disable=W0102 def __init__(
def __init__(self, command, log_file, env=os.environ.copy(), timeout=None): self,
command: str,
log_file: str,
env: Optional[dict] = None,
timeout: Optional[int] = None,
):
self.command = command self.command = command
self.log_file = log_file self._log_file_name = log_file
self.env = env self._log_file = None # type: Optional[TextIOWrapper]
self.env = env or os.environ.copy()
self._process = None # type: Optional[Popen] self._process = None # type: Optional[Popen]
self.timeout = timeout self.timeout = timeout
def _check_timeout(self): def _check_timeout(self) -> None:
if self.timeout is None:
return
sleep(self.timeout) sleep(self.timeout)
while self.process.poll() is None: while self.process.poll() is None:
logging.warning( logging.warning(
@ -33,7 +42,7 @@ class TeePopen:
os.killpg(self.process.pid, 9) os.killpg(self.process.pid, 9)
sleep(10) sleep(10)
def __enter__(self): def __enter__(self) -> "TeePopen":
self.process = Popen( self.process = Popen(
self.command, self.command,
shell=True, shell=True,
@ -44,25 +53,21 @@ class TeePopen:
stdout=PIPE, stdout=PIPE,
bufsize=1, bufsize=1,
) )
self.log_file = open(self.log_file, "w", encoding="utf-8")
if self.timeout is not None and self.timeout > 0: if self.timeout is not None and self.timeout > 0:
t = Thread(target=self._check_timeout) t = Thread(target=self._check_timeout)
t.daemon = True # does not block the program from exit t.daemon = True # does not block the program from exit
t.start() t.start()
return self return self
def __exit__(self, t, value, traceback): def __exit__(self, exc_type, exc_value, traceback):
for line in self.process.stdout: # type: ignore self.wait()
sys.stdout.write(line)
self.log_file.write(line)
self.process.wait()
self.log_file.close() self.log_file.close()
def wait(self): def wait(self):
for line in self.process.stdout: # type: ignore if self.process.stdout is not None:
sys.stdout.write(line) for line in self.process.stdout:
self.log_file.write(line) sys.stdout.write(line)
self.log_file.write(line)
return self.process.wait() return self.process.wait()
@ -75,3 +80,9 @@ class TeePopen:
@process.setter @process.setter
def process(self, process: Popen) -> None: def process(self, process: Popen) -> None:
self._process = process self._process = process
@property
def log_file(self) -> TextIOWrapper:
if self._log_file is None:
self._log_file = open(self._log_file_name, "w", encoding="utf-8")
return self._log_file