Files
SIA/sia/entry/single_entry.py
2024-12-04 17:19:46 +01:00

122 lines
3.9 KiB
Python

from pathlib import Path
import subprocess
import xml.etree.ElementTree as ET
from typing import Optional
from . import Entry
class SingleEntry(Entry):
"""
Entry type for one-time script executions.
"""
default_timeout = 1
default_limit = 1024
def __init__(
self,
id: str,
work_dir: Path,
script: str,
timeout: Optional[float] = None,
limit: Optional[int] = None,
):
"""
Initialize a new single shot entry.
Args:
id: Unique identifier for this entry
script: The script/command to execute
timeout: Maximum time to wait for script execution
limit: Maximum number of characters to capture from stdout/stderr
"""
super().__init__(id)
self.work_dir = work_dir
self.script = script
self.timeout = timeout
self.limit = limit
self.stdout: str = ""
self.stderr: str = ""
self.exit_code: Optional[int] = None
self.executed: bool = False
self.timed_out: bool = False
def reset(self) -> None:
"""Reset execution state to allow running again."""
self.executed = False
self.timed_out = False
self.stdout = ""
self.stderr = ""
self.exit_code = None
self.notify_change()
def update(self) -> None:
"""
Execute the script if not already executed.
Captures stdout, stderr and exit code.
"""
if self.executed:
return
self.executed = True
try:
process = subprocess.run(
self.script,
cwd=self.work_dir,
timeout=(self.timeout or self.default_timeout),
shell=True,
capture_output=True,
text=True
)
self.stdout = process.stdout
self.stderr = process.stderr
self.exit_code = process.returncode
except subprocess.TimeoutExpired as e:
self.timed_out = True
self.notify_change()
def generate_context(self) -> ET.Element:
"""
Generate an XML Element representing this single shot entry.
Returns:
ET.Element: XML element containing the entry's data
"""
element = ET.Element("single", {"id": self.id})
if self.timeout:
element.set("timeout", str(self.timeout))
element.text = self.script
if self.timed_out:
element.set("timed_out", "true")
elif self.executed:
element.set("exit_code", str(self.exit_code))
if self.limit:
element.set("limit", str(self.limit))
if len(self.stdout) > (self.limit or self.default_limit):
element.set("stdout_truncated", "true")
element.set("stdout_length", str(len(self.stdout)))
stdout_elem = ET.SubElement(element, "stdout")
stdout_elem.text = self.stdout[:(self.limit or self.default_limit)]
if len(self.stderr) > (self.limit or self.default_limit):
element.set("stderr_truncated", "true")
element.set("stderr_length", str(len(self.stderr)))
stderr_elem = ET.SubElement(element, "stderr")
stderr_elem.text = self.stderr[:(self.limit or self.default_limit)]
return element
def serialize(self):
"""
Collect all public attributes of the entry into a dictionary.
"""
return {
"type": "single",
"id": self.id,
"work_dir": str(self.work_dir),
"script": self.script,
"timeout": self.timeout,
"limit": self.limit,
"stdout": self.stdout,
"stderr": self.stderr,
"exit_code": self.exit_code,
"executed": self.executed,
"timed_out": self.timed_out
}