engine.py

"""
A 3dequalizer engine for SGTK.

"""
import os
import re
import logging
import shutil
import traceback

import tde4
from vl_sdv import rot3d, mat3d, VL_APPLY_ZXY

import sgtk
from sgtk.platform import Engine
from sgtk.util.filesystem import ensure_folder_exists

_HEARTBEAT_INTERVAL_MS = 50
_MENU_FOLDER_PATH = os.environ["TK_3DE4_MENU_FOLDER_PATH"]


class TDEqualizerEngine(Engine):
    def __init__(self, *args, **kwargs):
        self._current_file = tde4.getProjectPath()
        Engine.__init__(self, *args, **kwargs)

    def _heartbeat(self):
        from sgtk.platform.qt import QtCore

        # Keep Qt alive
        QtCore.QCoreApplication.processEvents()
        # check for open file change
        cur_file = tde4.getProjectPath()
        if self._current_file != cur_file:
            if cur_file:
                new_context = self.sgtk.context_from_path(cur_file, self.context)
                if new_context != self.context:
                    sgtk.platform.change_context(new_context)
            self._current_file = cur_file

    def pre_app_init(self):
        from sgtk.platform.qt import QtCore, QtGui

        if not QtCore.QCoreApplication.instance():
            # WARNING: need to keep a python reference to
            # the qt app, or python will destroy it and
            # ruin everything
            self._qt_app = QtGui.QApplication([])
            self._initialize_dark_look_and_feel()
            tde4.setTimerCallbackFunction(
                "sgtk.platform.current_engine()._heartbeat",
                _HEARTBEAT_INTERVAL_MS,
            )

    def post_app_init(self):
        self.create_shotgun_menu()

    def post_context_change(self, old_context, new_context):
        self.create_shotgun_menu()

    def destroy_engine(self):
        self.logger.debug("%s: Destroying...", self)
        self._cleanup_custom_scripts_dir_path()

    @property
    def context_change_allowed(self):
        return True

    @property
    def host_info(self):
        host_info = dict(name="3DEqualizer", version="unknown")
        try:
            host_info["name"], host_info["version"] = re.match(
                r"^([^\s]+)\s+(.*)$", tde4.get3DEVersion()
            ).groups()
        except Exception:
            # Fallback to initialized above
            pass

        return host_info

    def create_shotgun_menu(self):
        if self.has_ui:
            from sgtk.platform.qt import QtCore, QtGui

            self.logger.info("Creating ShotGrid menu...")

            self._cleanup_custom_scripts_dir_path()

            # Get temp folder path and create it if needed.
            ensure_folder_exists(_MENU_FOLDER_PATH)

            for i, (name, _) in enumerate(self.commands.items()):
                display_name = name
                # Very strange glitch
                if display_name == "Export...":
                    display_name = "Export ..."
                script_path = os.path.join(_MENU_FOLDER_PATH, "{:04d}.py".format(i))
                f = open(script_path, "w")
                f.write(
                    "\n".join(
                        (
                            "# 3DE4.script.name: {}".format(display_name),
                            "# 3DE4.script.gui:	Main Window::ShotGrid",
                            "if __name__ == '__main__':",
                            "   import sgtk",
                            "   sgtk.platform.current_engine().commands[{}]['callback']()".format(
                                repr(name)
                            ),
                        )
                    )
                )
                f.close()

            QtCore.QTimer.singleShot(0, tde4.rescanPythonDirs)

            self.logger.info("ShotGrid menu created.")

            return True
        return False

    def _cleanup_custom_scripts_dir_path(self):
        try:
            shutil.rmtree(_MENU_FOLDER_PATH)
        except OSError as e:
            self.logger.debug(traceback.format_exc())

    @property
    def has_ui(self):
        return True

    ##########################################################################################
    # logging

    def _emit_log_message(self, handler, record):
        if record.levelno < logging.INFO:
            formatter = logging.Formatter("Debug: ShotGrid %(basename)s: %(message)s")
        else:
            formatter = logging.Formatter("ShotGrid %(basename)s: %(message)s")
        msg = formatter.format(record)
        print(msg)

    def _create_dialog(self, title, bundle, widget, parent):
        from sgtk.platform.qt import QtCore

        dialog = super(TDEqualizerEngine, self)._create_dialog(
            title, bundle, widget, parent
        )
        dialog.raise_()
        dialog.activateWindow()
        return dialog

    #############
    # custom api

    @property
    def api(self):
        return self.import_module("tk_3dequalizer").api

    def get_scene_rotation_3d_zxy(self):
        return rot3d(mat3d(tde4.getSceneRotation3D())).angles(VL_APPLY_ZXY)

    def iter_all_cameras(self):
        return self.api.TDECamera.iter_all()

    def iter_selected_cameras(self):
        return self.api.TDECamera.iter_selected()

    def iter_all_point_groups(self):
        return self.api.TDEPointGroup.iter_all()

    def iter_selected_point_groups(self):
        return self.api.TDEPointGroup.iter_selected()