Refactor toolsets patch to patch package

lib/filetools.py

@@ -9,6 +9,9 @@
 from wlf.pathtools import make_path_finder
 
 ROOT = os.path.abspath(os.path.dirname(__file__))
+module_path = make_path_finder(__file__)  # pylint: disable = invalid-name
+plugin_folder_path = make_path_finder(  # pylint: disable = invalid-name
+    module_path())
 
 
 def path(*other):

lib/gui/menu.py

@@ -182,36 +182,6 @@ def _add_menu(menu, parent=nuke.menu("Nuke")):
     create_menu_by_dir(m, os.path.abspath(
         os.path.join(__file__, _plugin_path)))
 
-    # Enhance 'ToolSets' menu.
-    def _create_shared_toolsets():
-        if not nuke.selectedNodes():
-            nuke.message(utf8('未选中任何节点,不能创建工具集'))
-            return
-        filename = nuke.getInput('ToolSet name')
-        if filename:
-            nuke.createToolset(filename=os.path.join(
-                'Shared', filename), rootPath=RESOURCE_DIR)
-        _refresh_toolsets_menu()
-
-    def _refresh_toolsets_menu():
-        """Extended nuke function.  """
-
-        nukescripts.toolsets._refreshToolsetsMenu()  # pylint: disable=protected-access
-        m = nuke.menu('Nodes').addMenu('ToolSets')
-        m.addCommand('刷新'.encode('utf-8'), _refresh_toolsets_menu)
-        m.addCommand(
-            '创建共享工具集'.encode('utf-8'), _create_shared_toolsets)
-        m.addCommand(
-            '打开共享工具集文件夹'.encode('utf-8'),
-            lambda: webbrowser.open(
-                os.path.join(RESOURCE_DIR, 'ToolSets/Shared')))
-
-    if not getattr(nukescripts.toolsets, '_refreshToolsetsMenu', None):
-        setattr(nukescripts.toolsets, '_refreshToolsetsMenu',
-                nukescripts.toolsets.refreshToolsetsMenu)
-    _refresh_toolsets_menu()
-    nukescripts.toolsets.refreshToolsetsMenu = _refresh_toolsets_menu
-
 
 def create_menu_by_dir(parent, dir_):
     """Create menus by given folder structrue."""

lib/menu.py

@@ -14,6 +14,7 @@ def main():
     import enable_later
     import asset
     import patch.precomp
+    import patch.toolsets
 
     def _setup_cgtw():
         import cgtwn
@@ -36,6 +37,7 @@ def _setup_cgtw():
         _setup_cgtw()
 
     patch.precomp.enable()
+    patch.toolsets.enable()
 
 
 if __name__ == '__main__':

lib/patch/core.py

@@ -0,0 +1,121 @@
+# -*- coding=UTF-8 -*-
+"""Core for function patch.  """
+
+from __future__ import (absolute_import, division, print_function,
+                        unicode_literals)
+
+import logging
+
+import six
+
+LOGGER = logging.getLogger(__name__)
+
+
+class BasePatch(object):
+    """Undoable function patcher.  """
+
+    target = None
+    _orig = None
+
+    @classmethod
+    def func(cls, *args, **kwargs):
+        """Patched function.  """
+
+        raise NotImplementedError()
+
+    @classmethod
+    def orig(cls, *args, **kwargs):
+        """Shortcut for calling original function.  """
+
+        return cls._orig[0](*args, **kwargs)
+
+    @classmethod
+    def enable(cls, is_strict=True):
+        """Enabled patch.
+
+            Args:
+                is_strict(bool): If `is_strict` is True,
+                    patch on not existed attr will raise a AttributeError.
+        """
+
+        assert isinstance(cls.target, six.text_type), type(cls.target)
+        if cls._orig:
+            LOGGER.warning('Patch already enabled: %s', cls.target)
+            return
+        # Save as tuple, so target function
+        # will not become a unbound method of `Patch` obejct.
+        cls._orig = (resolve(cls.target),)
+        name, target = derive_importpath(cls.target, is_strict)
+        # Need convert classmethod to a function before set it.
+        setattr(target, name,
+                lambda *args, **kwargs: cls.func(*args, **kwargs))  # pylint:disable = unnecessary-lambda
+
+    @classmethod
+    def disable(cls):
+        """Disable patch.  """
+
+        if not cls._orig:
+            LOGGER.warning('Patch is not enabled: %s', cls.target)
+            return
+        name, target = derive_importpath(cls.target, True)
+        setattr(target, name, cls._orig[0])
+        cls._orig = None
+
+
+# Module name resolve function from `pytest.monkeypatch`
+
+def resolve(name):
+    # pylint: disable=missing-docstring
+    # simplified from zope.dottedname
+    parts = name.split('.')
+
+    used = parts.pop(0)
+    found = __import__(used)
+    for part in parts:
+        used += '.' + part
+        try:
+            found = getattr(found, part)
+        except AttributeError:
+            pass
+        else:
+            continue
+        # we use explicit un-nesting of the handling block in order
+        # to avoid nested exceptions on python 3
+        try:
+            __import__(used)
+        except ImportError as ex:
+            # str is used for py2 vs py3
+            expected = str(ex).split()[-1]
+            if expected == used:
+                raise
+            else:
+                raise ImportError(
+                    'import error in %s: %s' % (used, ex)
+                )
+        found = annotated_getattr(found, part, used)
+    return found
+
+
+def annotated_getattr(obj, name, ann):
+    # pylint: disable=missing-docstring
+    try:
+        obj = getattr(obj, name)
+    except AttributeError:
+        raise AttributeError(
+            '%r object at %s has no attribute %r' % (
+                type(obj).__name__, ann, name
+            )
+        )
+    return obj
+
+
+def derive_importpath(import_path, raising):
+    # pylint: disable=missing-docstring
+    if not isinstance(import_path, six.string_types) or "." not in import_path:
+        raise TypeError("must be absolute import path string, not %r" %
+                        (import_path,))
+    module, attr = import_path.rsplit('.', 1)
+    target = resolve(module)
+    if raising:
+        annotated_getattr(target, attr, ann=module)
+    return attr, target

lib/patch/precomp.py

@@ -5,22 +5,21 @@
                         unicode_literals)
 
 import nuke
-import nukescripts  # pylint: disable=import-error
 
 from wlf.codectools import get_unicode as u
 from wlf.path import PurePath
 
-UNPATCH_FUNC_STACK = []
+from .core import BasePatch
 
 
-def patch_precomp_dialog():
-    """Enhance precomp creation.  """
+class Patch(BasePatch):
+    """Enhance precomp creation """
+    target = 'nukescripts.PrecompOptionsDialog.__init__'
 
-    orig = nukescripts.PrecompOptionsDialog.__init__
-
-    def _func(self):
-
-        orig(self)
+    @classmethod
+    def func(cls, *args, **kwargs):
+        self = args[0]
+        cls.orig(*args, **kwargs)
 
         self.precompName = nuke.String_Knob(
             'precomp_name'.encode('utf-8'),
@@ -34,46 +33,33 @@ def _func(self):
         self.origNodes.setLabel('原节点'.encode('utf-8'))
 
         _knob_changed(self, self.precompName)
-        nuke.addKnobChanged(lambda self: _knob_changed(self, nuke.thisKnob()),
-                            args=self,
-                            nodeClass='PanelNode',
-                            node=self._PythonPanel__node)  # pylint: disable=protected-access
-
-    def _knob_changed(self, knob):
-        if knob is not self.precompName:
-            return
-
-        rootpath = PurePath(u(nuke.value('root.name')))
-        name = u(knob.value()) or 'precomp1'
-        self.scriptPath.setValue(
-            (
-                rootpath.parent /
-                ''.join([rootpath.stem]
-                        + ['.{}'.format(name)]
-                        + rootpath.suffixes)
-            ).as_posix().encode('utf-8'))
-        self.renderPath.setValue(
-            'precomp/{0}/{0}.%04d.exr'.format(
-                ''.join([rootpath.stem]
-                        + ['.{}'.format(name)]
-                        + [i for i in rootpath.suffixes if i != '.nk'])
-            ).encode('utf-8'))
-
-    nukescripts.PrecompOptionsDialog.__init__ = _func
-    UNPATCH_FUNC_STACK.append(lambda: setattr(
-        nukescripts.PrecompOptionsDialog, '__init__', orig))
-
-
-def enable():
-    """Enable patch. """
-    patch_precomp_dialog()
-
-
-def disable():
-    """Disable patch. """
-
-    while True:
-        try:
-            UNPATCH_FUNC_STACK.pop()()
-        except IndexError:
-            return
+        nuke.addKnobChanged(
+            lambda self: _knob_changed(self, nuke.thisKnob()),
+            args=self,
+            nodeClass='PanelNode',
+            node=self._PythonPanel__node)  # pylint: disable=protected-access
+
+
+def _knob_changed(self, knob):
+    if knob is not self.precompName:
+        return
+
+    rootpath = PurePath(u(nuke.value('root.name')))
+    name = u(knob.value()) or 'precomp1'
+    self.scriptPath.setValue(
+        (
+            rootpath.parent /
+            ''.join([rootpath.stem]
+                    + ['.{}'.format(name)]
+                    + rootpath.suffixes)
+        ).as_posix().encode('utf-8'))
+    self.renderPath.setValue(
+        'precomp/{0}/{0}.%04d.exr'.format(
+            ''.join([rootpath.stem]
+                    + ['.{}'.format(name)]
+                    + [i for i in rootpath.suffixes if i != '.nk'])
+        ).encode('utf-8'))
+
+
+enable = Patch.enable  # pylint: disable=invalid-name
+disable = Patch.disable  # pylint: disable=invalid-name

lib/patch/toolsets.py

@@ -0,0 +1,54 @@
+# -*- coding=UTF-8 -*-
+"""Patch nukescript toolset functions.  """
+
+from __future__ import (absolute_import, division, print_function,
+                        unicode_literals)
+
+import webbrowser
+
+import nuke
+import nukescripts  # pylint: disable=import-error
+
+import filetools
+
+from .core import BasePatch
+
+
+class Patch(BasePatch):
+    """Enhance toolsets menu.  """
+
+    target = 'nukescripts.toolsets.refreshToolsetsMenu'
+
+    @classmethod
+    def func(cls, *args, **kwargs):
+
+        cls.orig(*args, **kwargs)
+        m = nuke.menu('Nodes').addMenu('ToolSets')
+        m.addCommand('刷新'.encode('utf-8'), cls.func)
+        m.addCommand(
+            '创建共享工具集'.encode('utf-8'), _create_shared_toolsets)
+        m.addCommand(
+            '打开共享工具集文件夹'.encode('utf-8'),
+            lambda: webbrowser.open(
+                filetools.plugin_folder_path('ToolSets', 'Shared')))
+
+    @classmethod
+    def enable(cls, is_strict=True):
+        super(Patch, cls).enable(is_strict)
+        cls.func()
+
+
+def _create_shared_toolsets():
+    if not nuke.selectedNodes():
+        nuke.message('未选中任何节点,不能创建工具集'.encode('utf-8'))
+        return
+    toolset_name = nuke.getInput('ToolSet name')
+    if toolset_name:
+        nuke.createToolset(
+            filename='Shared/{}'.format(toolset_name),
+            rootPath=filetools.plugin_folder_path())
+    nukescripts.toolsets.refreshToolsetsMenu()
+
+
+enable = Patch.enable  # pylint: disable=invalid-name
+disable = Patch.disable  # pylint: disable=invalid-name

tests/test_filetools.py

@@ -0,0 +1,15 @@
+# -*- coding=UTF-8 -*-
+"""Test `filetools` module.  """
+
+from __future__ import (absolute_import, division, print_function,
+                        unicode_literals)
+
+import os
+
+import filetools
+from wlf.codectools import get_encoded as e
+
+
+def test_plugin_folder_path():
+    path = filetools.plugin_folder_path('ToolSets')
+    assert os.path.exists(e(path)), path