Refactor process title handling in RotaBuilder for improved clarity and functionality
This commit is contained in:
+1
-1
@@ -70,7 +70,7 @@ def main(
|
||||
Rota.constraint_options["max_weekend_frequency"] = 3
|
||||
Rota.constraint_options["hard_constrain_pair_separation"] = True
|
||||
|
||||
Rota.constraint_options["max_days_per_week_block"] = [(5, 2), (6, 3), (8,4)]
|
||||
Rota.constraint_options["max_days_per_week_block"] = [(8,4)]
|
||||
|
||||
#wr = [
|
||||
# WorkerRequirement(
|
||||
|
||||
@@ -38,6 +38,49 @@ from loguru import logger
|
||||
|
||||
logger.add("rota.log", rotation="1 MB", level="DEBUG")
|
||||
|
||||
import ctypes
|
||||
import platform
|
||||
|
||||
# Optional dependency: prefer setproctitle if available, fall back to libc on macOS
|
||||
try:
|
||||
import setproctitle # type: ignore
|
||||
except Exception:
|
||||
setproctitle = None
|
||||
|
||||
|
||||
def _get_proc_title() -> Optional[str]:
|
||||
try:
|
||||
if setproctitle is not None:
|
||||
return setproctitle.getproctitle()
|
||||
if platform.system() == "Darwin":
|
||||
libc = ctypes.CDLL("libc.dylib")
|
||||
libc.getprogname.restype = ctypes.c_char_p
|
||||
res = libc.getprogname()
|
||||
if res:
|
||||
return res.decode()
|
||||
except Exception:
|
||||
logger.debug("Could not read current process title")
|
||||
return None
|
||||
|
||||
|
||||
def _set_proc_title(title: str) -> bool:
|
||||
try:
|
||||
if setproctitle is not None:
|
||||
setproctitle.setproctitle(title)
|
||||
return True
|
||||
if platform.system() == "Darwin":
|
||||
libc = ctypes.CDLL("libc.dylib")
|
||||
# setprogname expects a C string
|
||||
try:
|
||||
libc.setprogname(ctypes.c_char_p(title.encode()))
|
||||
return True
|
||||
except Exception:
|
||||
# some platforms may not support setprogname via ctypes
|
||||
pass
|
||||
except Exception:
|
||||
logger.debug("Failed to set process title to %s", title)
|
||||
return False
|
||||
|
||||
|
||||
def _safe_constraint_info(obj):
|
||||
"""Return a tuple (type_name, safe_dump) for logging without calling repr()
|
||||
@@ -618,10 +661,17 @@ class RotaBuilder(object):
|
||||
allow_force_assignment_with_leave_conflict: bool = False,
|
||||
SHIFT_BOUNDS=SHIFT_BOUNDS,
|
||||
name: str = "",
|
||||
set_process_title: bool = True,
|
||||
bank_holidays: dict[datetime.date, str] = bank_holiday_map,
|
||||
constraint_options: RotaConstraintOptions | dict | None = None,
|
||||
):
|
||||
self.name = name
|
||||
self.set_process_title = set_process_title
|
||||
# record original process title so we can restore it after the run
|
||||
try:
|
||||
self._original_proc_title = _get_proc_title()
|
||||
except Exception:
|
||||
self._original_proc_title = None
|
||||
|
||||
console.print(
|
||||
Panel(
|
||||
@@ -909,6 +959,17 @@ class RotaBuilder(object):
|
||||
export_with_timestamp=False,
|
||||
solver="appsi_highs",
|
||||
):
|
||||
# Optionally set an informative process title so OS tools show which rota is running
|
||||
self._proc_title_changed = False
|
||||
if getattr(self, "set_process_title", False):
|
||||
title = f"rota-gen:{self.name}" if self.name else "rota-gen"
|
||||
try:
|
||||
if _set_proc_title(title):
|
||||
self._proc_title_changed = True
|
||||
logger.debug("Process title set to %s", title)
|
||||
except Exception:
|
||||
logger.debug("Could not set process title")
|
||||
|
||||
self.run_start_time = datetime.datetime.now()
|
||||
|
||||
self.build_shifts()
|
||||
@@ -925,6 +986,17 @@ class RotaBuilder(object):
|
||||
|
||||
self.run_end_time = datetime.datetime.now()
|
||||
|
||||
# Restore the original process title if we changed it
|
||||
if getattr(self, "_proc_title_changed", False):
|
||||
try:
|
||||
if getattr(self, "_original_proc_title", None) is not None:
|
||||
_set_proc_title(self._original_proc_title)
|
||||
else:
|
||||
# best-effort restore to a generic python name
|
||||
_set_proc_title("python")
|
||||
except Exception:
|
||||
logger.debug("Could not restore original process title")
|
||||
|
||||
if export:
|
||||
self.export_rota_to_html(self.name, timestamp_filename=export_with_timestamp)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user