Merge pull request #781 from ynput/feature/AY-6021_Addons-initialization-for-child-processes

Addons: Preparation steps for child processes

client/ayon_core/addon/__init__.py

@@ -9,11 +9,18 @@ from .interfaces import (
 )
 
 from .base import (
+    ProcessPreparationError,
+    ProcessContext,
     AYONAddon,
     AddonsManager,
     load_addons,
 )
 
+from .utils import (
+    ensure_addons_are_process_context_ready,
+    ensure_addons_are_process_ready,
+)
+
 
 __all__ = (
     "click_wrap",
@@ -24,7 +31,12 @@ __all__ = (
     "ITrayService",
     "IHostAddon",
 
+    "ProcessPreparationError",
+    "ProcessContext",
     "AYONAddon",
     "AddonsManager",
     "load_addons",
+
+    "ensure_addons_are_process_context_ready",
+    "ensure_addons_are_process_ready",
 )

client/ayon_core/addon/base.py

@@ -10,6 +10,7 @@ import threading
 import collections
 from uuid import uuid4
 from abc import ABC, abstractmethod
+from typing import Optional
 
 import appdirs
 import ayon_api
@@ -64,6 +65,56 @@ MOVED_ADDON_MILESTONE_VERSIONS = {
 }
 
 
+class ProcessPreparationError(Exception):
+    """Exception that can be used when process preparation failed.
+
+    The message is shown to user (either as UI dialog or printed). If
+        different error is raised a "generic" error message is shown to user
+        with option to copy error message to clipboard.
+
+    """
+    pass
+
+
+class ProcessContext:
+    """Context of child process.
+
+    Notes:
+        This class is used to pass context to child process. It can be used
+            to use different behavior of addon based on information in
+            the context.
+        The context can be enhanced in future versions.
+
+    Args:
+        addon_name (Optional[str]): Addon name which triggered process.
+        addon_version (Optional[str]): Addon version which triggered process.
+        project_name (Optional[str]): Project name. Can be filled in case
+            process is triggered for specific project. Some addons can have
+            different behavior based on project.
+        headless (Optional[bool]): Is process running in headless mode.
+
+    """
+    def __init__(
+        self,
+        addon_name: Optional[str] = None,
+        addon_version: Optional[str] = None,
+        project_name: Optional[str] = None,
+        headless: Optional[bool] = None,
+        **kwargs,
+    ):
+        if headless is None:
+            # TODO use lib function to get headless mode
+            headless = os.getenv("AYON_HEADLESS_MODE") == "1"
+        self.addon_name: Optional[str] = addon_name
+        self.addon_version: Optional[str] = addon_version
+        self.project_name: Optional[str] = project_name
+        self.headless: bool = headless
+
+        if kwargs:
+            unknown_keys = ", ".join([f'"{key}"' for key in kwargs.keys()])
+            print(f"Unknown keys in ProcessContext: {unknown_keys}")
+
+
 # Inherit from `object` for Python 2 hosts
 class _ModuleClass(object):
     """Fake module class for storing AYON addons.
@@ -584,7 +635,29 @@ class AYONAddon(ABC):
         Args:
             enabled_addons (list[AYONAddon]): Addons that are enabled.
         """
+        pass
 
+    def ensure_is_process_ready(
+        self, process_context: ProcessContext
+    ):
+        """Make sure addon is prepared for a process.
+
+        This method is called when some action makes sure that addon has set
+        necessary data. For example if user should be logged in
+        and filled credentials in environment variables this method should
+        ask user for credentials.
+
+        Implementation of this method is optional.
+
+        Note:
+            The logic can be similar to logic in tray, but tray does not require
+                to be logged in.
+
+        Args:
+            process_context (ProcessContext): Context of child
+                process.
+
+        """
         pass
 
     def get_global_environments(self):

client/ayon_core/addon/ui/process_ready_error.py

@@ -0,0 +1,132 @@
+import sys
+import json
+from typing import Optional
+
+from qtpy import QtWidgets, QtCore
+
+from ayon_core.style import load_stylesheet
+from ayon_core.tools.utils import get_ayon_qt_app
+
+
+class DetailDialog(QtWidgets.QDialog):
+    def __init__(self, detail, parent):
+        super().__init__(parent)
+
+        self.setWindowTitle("Detail")
+
+        detail_input = QtWidgets.QPlainTextEdit(self)
+        detail_input.setPlainText(detail)
+        detail_input.setReadOnly(True)
+
+        layout = QtWidgets.QVBoxLayout(self)
+        layout.addWidget(detail_input, 1)
+
+    def showEvent(self, event):
+        self.resize(600, 400)
+        super().showEvent(event)
+
+
+class ErrorDialog(QtWidgets.QDialog):
+    def __init__(
+        self,
+        message: str,
+        detail: Optional[str],
+        parent: Optional[QtWidgets.QWidget] = None
+    ):
+        super().__init__(parent)
+
+        self.setWindowTitle("Preparation failed")
+        self.setWindowFlags(
+            self.windowFlags() & ~QtCore.Qt.WindowContextHelpButtonHint
+        )
+
+        message_label = QtWidgets.QLabel(self)
+
+        detail_wrapper = QtWidgets.QWidget(self)
+
+        detail_label = QtWidgets.QLabel(detail_wrapper)
+
+        detail_layout = QtWidgets.QVBoxLayout(detail_wrapper)
+        detail_layout.setContentsMargins(0, 0, 0, 0)
+        detail_layout.addWidget(detail_label)
+
+        btns_wrapper = QtWidgets.QWidget(self)
+
+        copy_detail_btn = QtWidgets.QPushButton("Copy detail", btns_wrapper)
+        show_detail_btn = QtWidgets.QPushButton("Show detail", btns_wrapper)
+        confirm_btn = QtWidgets.QPushButton("Close", btns_wrapper)
+
+        btns_layout = QtWidgets.QHBoxLayout(btns_wrapper)
+        btns_layout.setContentsMargins(0, 0, 0, 0)
+        btns_layout.addWidget(copy_detail_btn, 0)
+        btns_layout.addWidget(show_detail_btn, 0)
+        btns_layout.addStretch(1)
+        btns_layout.addWidget(confirm_btn, 0)
+
+        layout = QtWidgets.QVBoxLayout(self)
+        layout.addWidget(message_label, 0)
+        layout.addWidget(detail_wrapper, 1)
+        layout.addWidget(btns_wrapper, 0)
+
+        copy_detail_btn.clicked.connect(self._on_copy_clicked)
+        show_detail_btn.clicked.connect(self._on_show_detail_clicked)
+        confirm_btn.clicked.connect(self._on_confirm_clicked)
+
+        self._message_label = message_label
+        self._detail_wrapper = detail_wrapper
+        self._detail_label = detail_label
+
+        self._copy_detail_btn = copy_detail_btn
+        self._show_detail_btn = show_detail_btn
+        self._confirm_btn = confirm_btn
+
+        self._detail_dialog = None
+
+        self._detail = detail
+
+        self.set_message(message, detail)
+
+    def showEvent(self, event):
+        self.setStyleSheet(load_stylesheet())
+        self.resize(320, 140)
+        super().showEvent(event)
+
+    def set_message(self, message, detail):
+        self._message_label.setText(message)
+        self._detail = detail
+
+        for widget in (
+            self._copy_detail_btn,
+            self._show_detail_btn,
+        ):
+            widget.setVisible(bool(detail))
+
+    def _on_copy_clicked(self):
+        if self._detail:
+            clipboard = QtWidgets.QApplication.clipboard()
+            clipboard.setText(self._detail)
+
+    def _on_show_detail_clicked(self):
+        if self._detail_dialog is None:
+            self._detail_dialog = DetailDialog(self._detail, self)
+        self._detail_dialog.show()
+
+    def _on_confirm_clicked(self):
+        self.accept()
+
+
+def main():
+    json_path = sys.argv[-1]
+    with open(json_path, "r") as stream:
+        data = json.load(stream)
+
+    message = data["message"]
+    detail = data["detail"]
+    app = get_ayon_qt_app()
+    dialog = ErrorDialog(message, detail)
+    dialog.show()
+    app.exec_()
+
+
+if __name__ == "__main__":
+    main()

client/ayon_core/addon/utils.py

@@ -0,0 +1,184 @@
+import os
+import sys
+import contextlib
+import tempfile
+import json
+import traceback
+from io import StringIO
+from typing import Optional
+
+from ayon_core.lib import run_ayon_launcher_process
+
+from .base import AddonsManager, ProcessContext, ProcessPreparationError
+
+
+def _handle_error(
+    process_context: ProcessContext,
+    message: str,
+    detail: Optional[str],
+):
+    """Handle error in process ready preparation.
+
+    Shows UI to inform user about the error, or prints the message
+        to stdout if running in headless mode.
+
+    Todos:
+        Make this functionality with the dialog as unified function, so it can
+            be used elsewhere.
+
+    Args:
+        process_context (ProcessContext): The context in which the
+            error occurred.
+        message (str): The message to show.
+        detail (Optional[str]): The detail message to show (usually
+            traceback).
+
+    """
+    if process_context.headless:
+        if detail:
+            print(detail)
+        print(f"{10*'*'}\n{message}\n{10*'*'}")
+        return
+
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    script_path = os.path.join(current_dir, "ui", "process_ready_error.py")
+    with tempfile.NamedTemporaryFile("w", delete=False) as tmp:
+        tmp_path = tmp.name
+        json.dump(
+            {"message": message, "detail": detail},
+            tmp.file
+        )
+
+    try:
+        run_ayon_launcher_process(
+            "--skip-bootstrap",
+            script_path,
+            tmp_path,
+            add_sys_paths=True,
+            creationflags=0,
+        )
+
+    finally:
+        os.remove(tmp_path)
+
+
+def _start_tray():
+    from ayon_core.tools.tray import make_sure_tray_is_running
+
+    make_sure_tray_is_running()
+
+
+def ensure_addons_are_process_context_ready(
+    process_context: ProcessContext,
+    addons_manager: Optional[AddonsManager] = None,
+    exit_on_failure: bool = True,
+) -> Optional[Exception]:
+    """Ensure all enabled addons are ready to be used in the given context.
+
+    Call this method only in AYON launcher process and as first thing
+        to avoid possible clashes with preparation. For example 'QApplication'
+        should not be created.
+
+    Args:
+        process_context (ProcessContext): The context in which the
+            addons should be prepared.
+        addons_manager (Optional[AddonsManager]): The addons
+            manager to use. If not provided, a new one will be created.
+        exit_on_failure (bool, optional): If True, the process will exit
+            if an error occurs. Defaults to True.
+
+    Returns:
+        Optional[Exception]: The exception that occurred during the
+            preparation, if any.
+
+    """
+    if addons_manager is None:
+        addons_manager = AddonsManager()
+
+    exception = None
+    message = None
+    failed = False
+    use_detail = False
+    # Wrap the output in StringIO to capture it for details on fail
+    # - but in case stdout was invalid on start of process also store
+    #   the tracebacks
+    tracebacks = []
+    output = StringIO()
+    with contextlib.redirect_stdout(output):
+        with contextlib.redirect_stderr(output):
+            for addon in addons_manager.get_enabled_addons():
+                addon_failed = True
+                try:
+                    addon.ensure_is_process_ready(process_context)
+                    addon_failed = False
+                except ProcessPreparationError as exc:
+                    exception = exc
+                    message = str(exc)
+                    print(f"Addon preparation failed: '{addon.name}'")
+                    print(message)
+
+                except BaseException as exc:
+                    exception = exc
+                    use_detail = True
+                    message = "An unexpected error occurred."
+                    formatted_traceback = "".join(traceback.format_exception(
+                        *sys.exc_info()
+                    ))
+                    tracebacks.append(formatted_traceback)
+                    print(f"Addon preparation failed: '{addon.name}'")
+                    print(message)
+                    # Print the traceback so it is in the stdout
+                    print(formatted_traceback)
+
+                if addon_failed:
+                    failed = True
+                    break
+
+    output_str = output.getvalue()
+    # Print stdout/stderr to console as it was redirected
+    print(output_str)
+    if not failed:
+        if not process_context.headless:
+            _start_tray()
+        return None
+
+    detail = None
+    if use_detail:
+        # In case stdout was not captured, use the tracebacks as detail
+        if not output_str:
+            output_str = "\n".join(tracebacks)
+        detail = output_str
+
+    _handle_error(process_context, message, detail)
+    if not exit_on_failure:
+        return exception
+    sys.exit(1)
+
+
+def ensure_addons_are_process_ready(
+    addons_manager: Optional[AddonsManager] = None,
+    exit_on_failure: bool = True,
+    **kwargs,
+) -> Optional[Exception]:
+    """Ensure all enabled addons are ready to be used in the given context.
+
+    Call this method only in AYON launcher process and as first thing
+        to avoid possible clashes with preparation. For example 'QApplication'
+        should not be created.
+
+    Args:
+        addons_manager (Optional[AddonsManager]): The addons
+            manager to use. If not provided, a new one will be created.
+        exit_on_failure (bool, optional): If True, the process will exit
+            if an error occurs. Defaults to True.
+        kwargs: The keyword arguments to pass to the ProcessContext.
+
+    Returns:
+        Optional[Exception]: The exception that occurred during the
+            preparation, if any.
+
+    """
+    context: ProcessContext = ProcessContext(**kwargs)
+    return ensure_addons_are_process_context_ready(
+        context, addons_manager, exit_on_failure
+    )

client/ayon_core/lib/execute.py

@@ -179,7 +179,7 @@ def clean_envs_for_ayon_process(env=None):
     return env
 
 
-def run_ayon_launcher_process(*args, **kwargs):
+def run_ayon_launcher_process(*args, add_sys_paths=False, **kwargs):
     """Execute AYON process with passed arguments and wait.
 
     Wrapper for 'run_process' which prepends AYON executable arguments
@@ -209,6 +209,15 @@ def run_ayon_launcher_process(*args, **kwargs):
         # - fill more if you find more
         env = clean_envs_for_ayon_process(os.environ)
 
+    if add_sys_paths:
+        new_pythonpath = list(sys.path)
+        lookup_set = set(new_pythonpath)
+        for path in (env.get("PYTHONPATH") or "").split(os.pathsep):
+            if path and path not in lookup_set:
+                new_pythonpath.append(path)
+                lookup_set.add(path)
+        env["PYTHONPATH"] = os.pathsep.join(new_pythonpath)
+
     return run_subprocess(args, env=env, **kwargs)