Nodes: Refactor parent classes to support other nodetree types

ros-dorian · ros-dorian · commit 3da61ebf36a7 · 2022-09-03T17:38:59.000+02:00
diff --git a/mitsuba-blender/nodes/__init__.py b/mitsuba-blender/nodes/__init__.py
@@ -1,74 +1,13 @@
-from bpy.utils import register_class, unregister_class
-from nodeitems_utils import NodeCategory, NodeItem, register_node_categories, unregister_node_categories
-
-from . import (
-    sockets
-)
-
-class MitsubaNodeCategoryMaterial(NodeCategory):
-    @classmethod
-    def poll(cls, context):
-        return context.space_data.tree_type == 'mitsuba_material_nodes'
-
-mitsuba_material_node_categories = [
-    MitsubaNodeCategoryMaterial('MITSUBA_MATERIAL_BSDF', 'BSDFs', items=[
-        NodeItem('MitsubaNodeTwosidedBSDF', label='Twosided'),
-        NodeItem('MitsubaNodeDiffuseBSDF', label='Diffuse'),
-        NodeItem('MitsubaNodeDielectricBSDF', label='Dielectric'),
-        NodeItem('MitsubaNodeThinDielectricBSDF', label='Thin Dielectric'),
-        NodeItem('MitsubaNodeRoughDielectricBSDF', label='Rough Dielectric'),
-        NodeItem('MitsubaNodeConductorBSDF', label='Conductor'),
-        NodeItem('MitsubaNodeRoughConductorBSDF', label='Rough Conductor'),
-        NodeItem('MitsubaNodePlasticBSDF', label='Plastic'),
-        NodeItem('MitsubaNodeRoughPlasticBSDF', label='Rough Plastic'),
-        NodeItem('MitsubaNodeBumpMapBSDF', label='Bump Map'),
-        NodeItem('MitsubaNodeNormalMapBSDF', label='Normal Map'),
-        NodeItem('MitsubaNodeBlendBSDF', label='Blend'),
-        NodeItem('MitsubaNodeMaskBSDF', label='Opacity Mask'),
-        NodeItem('MitsubaNodeNullBSDF', label='Null'),
-        NodeItem('MitsubaNodePrincipledBSDF', label='Principled'),
-    ]),
-
-    MitsubaNodeCategoryMaterial('MITSUBA_MATERIAL_TEXTURE', 'Textures', items=[
-        NodeItem('MitsubaNodeBitmapTexture', label='Bitmap'),
-        NodeItem('MitsubaNodeCheckerboardTexture', label='Checkerboard'),
-    ]),
-
-    MitsubaNodeCategoryMaterial('MITSUBA_MATERIAL_OUTPUT', 'Output', items=[
-        NodeItem('MitsubaNodeOutputMaterial', label='Output'),
-    ]),
-
-    MitsubaNodeCategoryMaterial('MITSUBA_MATERIAL_TRANSFORM', 'Transforms', items=[
-        NodeItem('MitsubaNode2DTransform', label='Transform 2D'),
-    ]),
-]
-
-classes = (
-    sockets.MitsubaSocketBSDF,
-    sockets.MitsubaSocketColorTexture,
-    sockets.MitsubaSocketNormalMap,
-    sockets.MitsubaSocketFloatTextureNoDefault,
-    sockets.MitsubaSocketFloatTextureUnbounded,
-    sockets.MitsubaSocketFloatTextureBounded0to1,
-    sockets.MitsubaSocket2DTransform,
-)
-
 def register():
-    register_node_categories('MITSUBA_MATERIAL_TREE', mitsuba_material_node_categories)
-
-    from . import (materials, textures, transforms)
+    from . import (materials, textures, transforms, sockets)
+    sockets.register()
     transforms.register()
-    materials.register()
     textures.register()
-    for cls in classes:
-        register_class(cls)
+    materials.register()
 
 def unregister():
-    unregister_node_categories('MITSUBA_MATERIAL_TREE')
-
-    from . import (materials, textures, transforms)
+    from . import (materials, textures, transforms, sockets)
     materials.unregister()
     textures.unregister()
     transforms.unregister()
-    for cls in classes:
-        unregister_class(cls)
+    sockets.unregister()
diff --git a/mitsuba-blender/nodes/base.py b/mitsuba-blender/nodes/base.py
@@ -1,4 +1,80 @@
-import bpy
+from bpy.props import BoolProperty
+
+from ..utils.nodetree import get_output_nodes
+
+class MitsubaSocket:
+    '''
+    Base class for custom Mitsuba node sockets
+    '''
+    bl_label = ''
+
+    color = (1, 1, 1, 1)
+    slider = False
+
+    @classmethod
+    def is_input_connection_valid(cls, connected_socket):
+        for valid_cls in cls.valid_inputs:
+            if isinstance(connected_socket, valid_cls):
+                return True
+        return False
+
+    def get_linked_node(self):
+        if self.is_linked and len(self.links) > 0:
+            return self.links[0].from_node
+        return None
+
+    def has_valid_state(self):
+        # If this socket is linked, check that the connected
+        # socket is of a valid type.
+        if self.is_linked and len(self.links) > 0:
+            link = self.links[0]
+            if link and hasattr(self, 'valid_inputs'):
+                if not self.is_input_connection_valid(link.from_socket):
+                    return False
+        return True
+
+    def draw_prop(self, context, layout, node, text):
+        layout.prop(self, 'default_value', text=text, slider=self.slider)
+
+    def draw(self, context, layout, node, text):
+        if not self.has_valid_state():
+            layout.label(text='Wrong Input', icon='CANCEL')
+            return
+
+        has_default = hasattr(self, "default_value") and self.default_value is not None
+
+        if self.is_output or self.is_linked or not has_default:
+            layout.label(text=text)
+        else:
+            self.draw_prop(context, layout, node, text)
+
+    def draw_color(self, context, node):
+        return self.color
+
+    def to_default_dict(self, export_context):
+        '''
+        Export the default value of a socket in a form the Mitsuba can understand
+        (either a value or a dictionary for complex types).
+        '''
+        # Implement in subclasses
+        raise RuntimeError(f'{self.bl_idname} default conversion not implemented')
+
+    def to_dict(self, export_context):
+        '''
+        Convert this socket to a Mitsuba dictionary.
+        '''
+        if not self.has_valid_state():
+            export_context.log('Cannot export material: Invalid socket state', 'ERROR')
+            return None
+        linked_node = self.get_linked_node()
+        if linked_node is not None:
+            # If a socket is connected, convert the connected node to a Mitsuba dictionary
+            return linked_node.to_dict(export_context)
+        if hasattr(self, 'default_value'):
+            # If the socket is not connected and has a default value, convert it to a Mitsuba dictionary
+            return self.to_default_dict(export_context)
+        # Otherwise, this socket cannot be converted
+        return None
 
 class MitsubaNode:
     '''
@@ -32,3 +108,51 @@ def update(self):
     def to_dict(self, export_context):
         # To implement in the subclasses
         raise RuntimeError(f'{self.bl_idname} conversion not implemented.')
+
+class MitsubaNodeOutput(MitsubaNode):
+    '''
+    Shader node representing a Mitsuba material output
+    '''
+    def _update_active(output_node, context):
+        if not output_node.is_active:
+            output_node.is_active = True
+        output_node.disable_other_outputs()
+
+    is_active: BoolProperty(name='Active', default=True, update=_update_active)
+
+    def init(self, context):
+        super().init(context)
+        self.disable_other_outputs()
+
+    def draw_buttons(self, context, layout):
+        layout.prop(self, 'is_active')
+
+    def free(self):
+        if not self.is_active:
+            return
+
+        node_tree = self.id_data
+        if node_tree is None:
+            return
+        for node in get_output_nodes(node_tree):
+            if node != self:
+                node['is_active'] = True
+                break
+
+    def disable_other_outputs(self):
+        node_tree = self.id_data
+        if node_tree is None:
+            return
+        for node in get_output_nodes(node_tree):
+            if node != self:
+                node['is_active'] = False
+
+class MitsubaNodeTree:
+    '''
+    Base class for custom Mitsuba shader node trees
+    '''
+    bl_label = ''
+
+    @classmethod
+    def poll(cls, context):
+        return context.scene.render.engine == 'MITSUBA'
diff --git a/mitsuba-blender/nodes/materials/__init__.py b/mitsuba-blender/nodes/materials/__init__.py
@@ -1,9 +1,12 @@
-import bpy
 from bpy.utils import register_class, unregister_class
+from nodeitems_utils import register_node_categories, unregister_node_categories
+
 from . import (
     blend, bumpmap, nodetree, output, twosided, diffuse, dielectric, conductor, plastic, normalmap, mask, null, principled
 )
 
+from .nodetree import mitsuba_node_categories_material
+
 classes = (
     nodetree.MitsubaNodeTreeMaterial,
     output.MitsubaNodeOutputMaterial,
@@ -25,9 +28,13 @@
 )
 
 def register():
+    register_node_categories('MITSUBA_MATERIAL_TREE', mitsuba_node_categories_material)
+
     for cls in classes:
         register_class(cls)
 
 def unregister():
+    unregister_node_categories('MITSUBA_MATERIAL_TREE')
+
     for cls in classes:
         unregister_class(cls)
diff --git a/mitsuba-blender/nodes/materials/nodetree.py b/mitsuba-blender/nodes/materials/nodetree.py
@@ -1,17 +1,16 @@
 import bpy
+from nodeitems_utils import NodeCategory, NodeItem
 
-class MitsubaNodeTreeMaterial(bpy.types.NodeTree):
+from ..base import MitsubaNodeTree
+
+class MitsubaNodeTreeMaterial(bpy.types.NodeTree, MitsubaNodeTree):
     '''
     Custom Blender Node Tree for Mitsuba BSDFs
     '''
     bl_idname = 'mitsuba_material_nodes'
     bl_label = 'Mitsuba Material Editor'
     bl_icon = 'NODE_MATERIAL'
 
-    @classmethod
-    def poll(cls, context):
-        return context.scene.render.engine == 'MITSUBA'
-
     @classmethod
     def get_from_context(cls, context):
         '''
@@ -25,3 +24,41 @@ def get_from_context(cls, context):
                 if node_tree:
                     return node_tree, mat, mat
         return None, None, None
+
+class MitsubaNodeCategoryMaterial(NodeCategory):
+    @classmethod
+    def poll(cls, context):
+        return context.space_data.tree_type == MitsubaNodeTreeMaterial.bl_idname
+
+mitsuba_node_categories_material = [
+    MitsubaNodeCategoryMaterial('MITSUBA_MATERIAL_BSDF', 'BSDFs', items=[
+        NodeItem('MitsubaNodeTwosidedBSDF', label='Twosided'),
+        NodeItem('MitsubaNodeDiffuseBSDF', label='Diffuse'),
+        NodeItem('MitsubaNodeDielectricBSDF', label='Dielectric'),
+        NodeItem('MitsubaNodeThinDielectricBSDF', label='Thin Dielectric'),
+        NodeItem('MitsubaNodeRoughDielectricBSDF', label='Rough Dielectric'),
+        NodeItem('MitsubaNodeConductorBSDF', label='Conductor'),
+        NodeItem('MitsubaNodeRoughConductorBSDF', label='Rough Conductor'),
+        NodeItem('MitsubaNodePlasticBSDF', label='Plastic'),
+        NodeItem('MitsubaNodeRoughPlasticBSDF', label='Rough Plastic'),
+        NodeItem('MitsubaNodeBumpMapBSDF', label='Bump Map'),
+        NodeItem('MitsubaNodeNormalMapBSDF', label='Normal Map'),
+        NodeItem('MitsubaNodeBlendBSDF', label='Blend'),
+        NodeItem('MitsubaNodeMaskBSDF', label='Opacity Mask'),
+        NodeItem('MitsubaNodeNullBSDF', label='Null'),
+        NodeItem('MitsubaNodePrincipledBSDF', label='Principled'),
+    ]),
+
+    MitsubaNodeCategoryMaterial('MITSUBA_MATERIAL_TEXTURE', 'Textures', items=[
+        NodeItem('MitsubaNodeBitmapTexture', label='Bitmap'),
+        NodeItem('MitsubaNodeCheckerboardTexture', label='Checkerboard'),
+    ]),
+
+    MitsubaNodeCategoryMaterial('MITSUBA_MATERIAL_OUTPUT', 'Output', items=[
+        NodeItem('MitsubaNodeOutputMaterial', label='Output'),
+    ]),
+
+    MitsubaNodeCategoryMaterial('MITSUBA_MATERIAL_TRANSFORM', 'Transforms', items=[
+        NodeItem('MitsubaNode2DTransform', label='Transform 2D'),
+    ]),
+]
diff --git a/mitsuba-blender/nodes/materials/output.py b/mitsuba-blender/nodes/materials/output.py
@@ -1,36 +1,19 @@
 import bpy
-from bpy.props import BoolProperty
-from ..base import MitsubaNode
+from ..base import MitsubaNodeOutput
 
-class MitsubaNodeOutputMaterial(bpy.types.Node, MitsubaNode):
+class MitsubaNodeOutputMaterial(bpy.types.Node, MitsubaNodeOutput):
     '''
     Shader node representing a Mitsuba material output
     '''
     bl_idname = 'MitsubaNodeOutputMaterial'
     bl_label = 'Material Output'
 
-    def _update_active(output_node, context):
-        if not output_node.is_active:
-            output_node.is_active = True
-        output_node.disable_other_outputs()
-
-    is_active: BoolProperty(name='Active', default=True, update=_update_active)
-
     def init(self, context):
-        super().init(context)
         self.add_input('MitsubaSocketBSDF', 'BSDF')
-        self.disable_other_outputs()
+        super().init(context)
 
     def draw_buttons(self, context, layout):
-        layout.prop(self, 'is_active')
-
-    def disable_other_outputs(self):
-        node_tree = self.id_data
-        if node_tree is None:
-            return
-        for node in node_tree.nodes:
-            if getattr(node, 'bl_idname', None) == 'MitsubaNodeOutputMaterial' and node != self:
-                node['is_active'] = False
+        super().draw_buttons(context, layout)
 
     def to_dict(self, export_context):
         if not self.is_active:
diff --git a/mitsuba-blender/nodes/sockets.py b/mitsuba-blender/nodes/sockets.py
diff --git a/mitsuba-blender/utils/nodetree.py b/mitsuba-blender/utils/nodetree.py
