Refactor Standalone section shader (#2172)

* Refactor Standalone section shader

* Remove leftover code

src/blenderbim/scripts/standalone_section_shader.py

@@ -13,11 +13,9 @@ bl_info = {
 
 
 import bpy
-from bpy.types import Operator
+from bpy.types import Operator, PropertyGroup, Panel
-from bpy.props import FloatVectorProperty
 from mathutils import Vector, Matrix, Euler
 from math import radians
-from bpy.types import PropertyGroup
 from bpy.props import (
     PointerProperty,
     StringProperty,
@@ -30,62 +28,131 @@ from bpy.props import (
 )
 
 
-def update_section_color(self, context):
+class SectionCutawayManager:
-    section_node_group = bpy.data.node_groups.get("Section Override")
+    @staticmethod
-    if section_node_group is None:
+    def get_section_tree():
-        return
+        return bpy.data.node_groups.get("Section Override")
-    try:
-        emission_node = next(n for n in section_node_group.nodes if isinstance(n, bpy.types.ShaderNodeEmission))
-        emission_node.inputs[0].default_value = list(self.section_plane_colour) + [1]
-    except StopIteration:
-        pass
 
+    @staticmethod
+    def get_sections_map():
+        section_tree = SectionCutawayManager.get_section_tree()
+        return {n: n.object for n in section_tree.nodes if isinstance(n, bpy.types.ShaderNodeTexCoord)}
 
-class BIM_OT_add_section_plane(bpy.types.Operator):
+    @staticmethod
-    """Add a temporary empty object as a section cutaway. Cull all geometry rendering below the empty's local Z axis"""
+    def purge_all_section_data():
+        section_override_mat = bpy.data.materials.get("Section Override")
+        if section_override_mat:
+            bpy.data.materials.remove(section_override_mat)
+        for material in bpy.data.materials:
+            if not material.node_tree:
+                continue
+            override = material.node_tree.nodes.get("Section Override")
+            if not override:
+                continue
+            material.node_tree.links.new(
+                override.inputs[0].links[0].from_socket, override.outputs[0].links[0].to_socket
+            )
+            material.node_tree.nodes.remove(override)
+        bpy.data.node_groups.remove(SectionCutawayManager.get_section_tree())
+        bpy.data.node_groups.remove(bpy.data.node_groups.get("Section Compare"))
 
-    bl_idname = "bim.add_section_plane"
+    @staticmethod
-    bl_label = "Add Temporary Section Cutaway"
+    def clean():
-    bl_options = {"REGISTER", "UNDO"}
+        sections_map = SectionCutawayManager.get_sections_map()
+        sections_map_iterator = list(sections_map.items())
+        for tex_coord_node, section_obj in sections_map_iterator:
+            if section_obj is None or section_obj.users == 1:
+                SectionCutawayManager.unplug(tex_coord_node)
+                sections_map.pop(tex_coord_node)
+        if not any(o for o in sections_map.values() if o.users > 1):
+            SectionCutawayManager.purge_all_section_data()
 
-    def execute(self, context):
+    @staticmethod
-        obj = self.create_section_obj(context)
+    def unplug(tex_coords):
-        if not self.has_section_override_node():
+        section_tree = SectionCutawayManager.get_section_tree()
-            self.create_section_compare_node()
+        section_compare = tex_coords.outputs["Object"].links[0].to_node
-            self.create_section_override_node(obj, context)
+        if section_compare.inputs[0].links:
-        else:
+            previous_section_compare = section_compare.inputs[0].links[0].from_node
-            self.append_obj_to_section_override_node(obj)
+            next_section_compare = section_compare.outputs[0].links[0].to_node
-        self.add_default_material_if_none_exists(context)
+            section_tree.links.new(previous_section_compare.outputs[0], next_section_compare.inputs[0])
-        self.override_materials()
+            SectionCutawayManager.offset_previous_nodes(section_compare, offset_x=200)
-        return {"FINISHED"}
+        section_tree.nodes.remove(section_compare)
+        section_tree.nodes.remove(tex_coords)
 
-    def create_section_obj(self, context):
+    @staticmethod
+    def offset_previous_nodes(section_compare, offset_x=0, offset_y=0):
+        if section_compare.inputs[0].links:
+            previous_section_compare = section_compare.inputs[0].links[0].from_node
+            previous_section_compare.location += Vector((offset_x, offset_y))
+            if previous_section_compare.inputs[1].links:
+                previous_section_compare.inputs[1].links[0].from_node.location += Vector((offset_x, offset_y))
+            SectionCutawayManager.offset_previous_nodes(previous_section_compare, offset_x, offset_y)
+
+    @staticmethod
+    def override_materials():
+        override = SectionCutawayManager.get_section_tree()
+        for material in bpy.data.materials:
+            material.use_nodes = True
+            if material.node_tree.nodes.get(override.name):
+                continue
+            material.blend_method = "HASHED"
+            material.shadow_method = "HASHED"
+            material_output = next((n for n in material.node_tree.nodes if n.type == "OUTPUT_MATERIAL"), None)
+            if material_output is None:
+                continue
+            from_socket = material_output.inputs[0].links[0].from_socket
+            section_override = material.node_tree.nodes.new(type="ShaderNodeGroup")
+            section_override.name = "Section Override"
+            section_override.node_tree = override
+            material.node_tree.links.new(from_socket, section_override.inputs[0])
+            material.node_tree.links.new(section_override.outputs[0], material_output.inputs[0])
+
+    @staticmethod
+    def create_section_obj(from_obj, scene):
         section = bpy.data.objects.new("Section", None)
         section.empty_display_type = "SINGLE_ARROW"
         section.empty_display_size = 5
         section.show_in_front = True
-        if (
+        if from_obj and from_obj.select_get() and isinstance(from_obj.data, bpy.types.Camera):
-            context.active_object
+            section.matrix_world = from_obj.matrix_world @ Euler((radians(180.0), 0.0, 0.0), "XYZ").to_matrix().to_4x4()
-            and context.active_object.select_get()
-            and isinstance(context.active_object.data, bpy.types.Camera)
-        ):
-            section.matrix_world = (
-                context.active_object.matrix_world @ Euler((radians(180.0), 0.0, 0.0), "XYZ").to_matrix().to_4x4()
-            )
         else:
             section.rotation_euler = Euler((radians(180.0), 0.0, 0.0), "XYZ")
-            section.location = context.scene.cursor.location
+            section.location = scene.cursor.location
         collection = bpy.data.collections.get("Sections")
         if not collection:
             collection = bpy.data.collections.new("Sections")
-            context.scene.collection.children.link(collection)
+            scene.collection.children.link(collection)
         collection.objects.link(section)
         return section
 
-    def has_section_override_node(self):
+    @classmethod
-        return bpy.data.node_groups.get("Section Override")
+    def add_section_plane(cls, context):
+        obj = cls.create_section_obj(context.active_object, context.scene)
+        if cls.get_section_tree() is not None:
+            cls.append_obj_to_section_override_node(obj)
+        else:
+            cls.create_section_compare_node()
+            cls.create_section_override_node(obj, context)
+        cls.add_default_material_if_none_exists(context)
+        cls.override_materials()
 
-    def create_section_compare_node(self):
+    @classmethod
+    def remove_all_section_planes(cls):
+        objs = [o for o in cls.get_sections_map().values() if o is not None]
+        cls.remove_section_planes(objs)
+
+    @classmethod
+    def remove_section_planes(cls, objects):
+        sections_map = cls.get_sections_map()
+        for obj in objects:
+            tex_coords = next((n for n in sections_map if sections_map[n] == obj), None)
+            if tex_coords is not None:
+                cls.unplug(tex_coords)
+                sections_map.pop(tex_coords)
+                bpy.data.objects.remove(obj)
+
+    @staticmethod
+    def create_section_compare_node():
         group = bpy.data.node_groups.new("Section Compare", type="ShaderNodeTree")
         group_input = group.nodes.new(type="NodeGroupInput")
         group_input.location = 0, 50
@@ -112,7 +179,8 @@ class BIM_OT_add_section_plane(bpy.types.Operator):
         group.links.new(greater.outputs[0], multiply.inputs[1])
         group.links.new(multiply.outputs[0], group_output.inputs[""])
 
-    def create_section_override_node(self, obj, context):
+    @staticmethod
+    def create_section_override_node(obj, context):
         group = bpy.data.node_groups.new("Section Override", type="ShaderNodeTree")
         links = group.links
         nodes = group.nodes
@@ -121,45 +189,53 @@ class BIM_OT_add_section_plane(bpy.types.Operator):
         group_output = nodes.new(type="NodeGroupOutput")
         group_output.location = 600, 250
 
+        hide_mix = group.nodes.new(type="ShaderNodeMixShader")
+        hide_mix.name = "Hide Mix"
+        hide_mix.inputs[0].default_value = 0
+        hide_mix.location = group_output.location - Vector((200, 0))
+
         backfacing_mix = nodes.new(type="ShaderNodeMixShader")
-        backfacing_mix.location = group_output.location - Vector((400, 350))
+        backfacing_mix.location = group_output.location - Vector((600, 350))
 
         backfacing = nodes.new(type="ShaderNodeNewGeometry")
-        backfacing.location = backfacing_mix.location + Vector((-200, 200))
+        backfacing.location = backfacing_mix.location + Vector((-400, 200))
-        group_input.location = backfacing_mix.location - Vector((200, 50))
+        group_input.location = backfacing_mix.location - Vector((400, 50))
 
         emission = nodes.new(type="ShaderNodeEmission")
         emission.inputs[0].default_value = list(context.scene.SectionProperties.section_plane_colour) + [1]
-        emission.location = backfacing_mix.location - Vector((200, 150))
+        emission.location = backfacing_mix.location - Vector((400, 150))
 
         transparent = nodes.new(type="ShaderNodeBsdfTransparent")
-        transparent.location = group_output.location - Vector((400, 100))
+        transparent.location = group_output.location - Vector((600, 100))
 
         section_mix = group.nodes.new(type="ShaderNodeMixShader")
         section_mix.name = "Section Mix"
         section_mix.inputs[0].default_value = 1  # Directly pass input shader when there is no cutaway
-        section_mix.location = group_output.location - Vector((200, 0))
+        section_mix.location = group_output.location - Vector((400, 0))
 
         cut_obj = nodes.new(type="ShaderNodeTexCoord")
         cut_obj.object = obj
-        cut_obj.location = group_output.location - Vector((800, 150))
+        cut_obj.location = group_output.location - Vector((1000, 150))
 
         section_compare = nodes.new(type="ShaderNodeGroup")
         section_compare.node_tree = bpy.data.node_groups.get("Section Compare")
         section_compare.name = "Last Section Compare"
-        section_compare.location = group_output.location - Vector((600, 0))
+        section_compare.location = group_output.location - Vector((800, 0))
 
         links.new(cut_obj.outputs["Object"], section_compare.inputs[1])
         links.new(backfacing.outputs["Backfacing"], backfacing_mix.inputs[0])
         links.new(group_input.outputs[""], backfacing_mix.inputs[1])
         links.new(emission.outputs["Emission"], backfacing_mix.inputs[2])
         links.new(section_compare.outputs[0], section_mix.inputs[0])
-        links.new(transparent.outputs["BSDF"], section_mix.inputs[1])
+        links.new(transparent.outputs[0], section_mix.inputs[1])
-        links.new(backfacing_mix.outputs["Shader"], section_mix.inputs[2])
+        links.new(backfacing_mix.outputs[0], section_mix.inputs[2])
-        links.new(section_mix.outputs["Shader"], group_output.inputs[""])
+        links.new(section_mix.outputs[0], hide_mix.inputs[1])
+        links.new(group_input.outputs["Shader"], hide_mix.inputs[2])
+        links.new(hide_mix.outputs[0], group_output.inputs[""])
 
-    def append_obj_to_section_override_node(self, obj):
+    @staticmethod
-        group = bpy.data.node_groups.get("Section Override")
+    def append_obj_to_section_override_node(obj):
+        group = SectionCutawayManager.get_section_tree()
         try:
             last_section_node = next(
                 n
@@ -185,7 +261,8 @@ class BIM_OT_add_section_plane(bpy.types.Operator):
 
         section_compare.name = "Last Section Compare"
 
-    def add_default_material_if_none_exists(self, context):
+    @staticmethod
+    def add_default_material_if_none_exists(context):
         material = bpy.data.materials.get("Section Override")
         if not material:
             material = bpy.data.materials.new("Section Override")
@@ -208,92 +285,74 @@ class BIM_OT_add_section_plane(bpy.types.Operator):
                     else:
                         obj.data.materials.append(material)
 
-    def override_materials(self):
-        override = bpy.data.node_groups.get("Section Override")
-        for material in bpy.data.materials:
-            material.use_nodes = True
-            if material.node_tree.nodes.get("Section Override"):
-                continue
-            material.blend_method = "HASHED"
-            material.shadow_method = "HASHED"
-            material_output = self.get_node(material.node_tree.nodes, "OUTPUT_MATERIAL")
-            if not material_output:
-                continue
-            from_socket = material_output.inputs[0].links[0].from_socket
-            section_override = material.node_tree.nodes.new(type="ShaderNodeGroup")
-            section_override.name = "Section Override"
-            section_override.node_tree = override
-            material.node_tree.links.new(from_socket, section_override.inputs[0])
-            material.node_tree.links.new(section_override.outputs[0], material_output.inputs[0])
 
-    def get_node(self, nodes, node_type):
+class BIM_OT_section_plane_refresh(Operator):
-        for node in nodes:
+    bl_idname = "bim.section_plane_refresh"
-            if node.type == node_type:
+    bl_label = "Refresh Section Cutaway Shader"
-                return node
+    bl_options = {"REGISTER", "UNDO"}
+
+    def execute(self, context):
+        update_visibility(None, context)
+        SectionCutawayManager.add_default_material_if_none_exists(context)
+        SectionCutawayManager.clean()
+        return {"FINISHED"}
 
 
-class BIM_OT_remove_section_plane(bpy.types.Operator):
+class BIM_OT_section_plane_add(Operator):
-    """Remove selected section plane. No effect if executed on a regular object"""
+    """Add a temporary empty object as a section cutaway. Cull all geometry rendering below the empty's local Z axis"""
 
-    bl_idname = "bim.remove_section_plane"
+    bl_idname = "bim.section_plane_add"
+    bl_label = "Add Temporary Section Cutaway"
+    bl_options = {"REGISTER", "UNDO"}
+
+    def execute(self, context):
+        SectionCutawayManager.add_section_plane(context)
+        SectionCutawayManager.clean()
+        return {"FINISHED"}
+
+
+class BIM_OT_section_plane_remove(Operator):
+    bl_idname = "bim.section_plane_remove"
     bl_label = "Remove Temporary Section Cutaway"
     bl_options = {"REGISTER", "UNDO"}
+    bl_description = "Remove section planes. No effect if executed on a regular object"
+    remove_all: BoolProperty(default=False, options={"HIDDEN"})
 
     @classmethod
     def poll(cls, context):
-        return context.active_object and bpy.data.node_groups.get("Section Override")
+        return SectionCutawayManager.get_section_tree() is not None
 
     def execute(self, context):
-        name = context.active_object.name
+        if self.remove_all:
-        section_override = bpy.data.node_groups.get("Section Override")
+            SectionCutawayManager.remove_all_section_planes()
-        tex_coords = next(
+        else:
-            (
+            SectionCutawayManager.remove_section_planes(context.selected_objects)
-                n
+        SectionCutawayManager.clean()
-                for n in section_override.nodes
-                if isinstance(n, bpy.types.ShaderNodeTexCoord) and n.object.name == name
-            ),
-            None,
-        )
-        if tex_coords is not None:
-            section_compare = tex_coords.outputs["Object"].links[0].to_node
-            if section_compare.inputs[0].links:
-                previous_section_compare = section_compare.inputs[0].links[0].from_node
-                next_section_compare = section_compare.outputs[0].links[0].to_node
-                section_override.links.new(previous_section_compare.outputs[0], next_section_compare.inputs[0])
-                self.offset_previous_nodes(section_compare, offset_x=200)
-            section_override.nodes.remove(section_compare)
-            section_override.nodes.remove(tex_coords)
-            bpy.data.objects.remove(context.active_object)
-
         return {"FINISHED"}
 
-    def offset_previous_nodes(self, section_compare, offset_x=0, offset_y=0):
-        if section_compare.inputs[0].links:
-            previous_section_compare = section_compare.inputs[0].links[0].from_node
-            previous_section_compare.location += Vector((offset_x, offset_y))
-            if previous_section_compare.inputs[1].links:
-                previous_section_compare.inputs[1].links[0].from_node.location += Vector((offset_x, offset_y))
-            self.offset_previous_nodes(previous_section_compare, offset_x, offset_y)
 
-    def purge_all_section_data(self, context):
+def update_visibility(self, context):
-        bpy.data.materials.remove(bpy.data.materials.get("Section Override"))
+    tree = SectionCutawayManager.get_section_tree()
-        for material in bpy.data.materials:
+    if tree:
-            if not material.node_tree:
+        node = tree.nodes.get("Hide Mix")
-                continue
+        if node:
-            override = material.node_tree.nodes.get("Section Override")
+            node.inputs[0].default_value = int(context.scene.SectionProperties.hide)
-            if not override:
+
-                continue
+
-            material.node_tree.links.new(
+def update_section_color(self, context):
-                override.inputs[0].links[0].from_socket, override.outputs[0].links[0].to_socket
+    section_node_group = SectionCutawayManager.get_section_tree()
-            )
+    if section_node_group is None:
-            material.node_tree.nodes.remove(override)
+        return
-        bpy.data.node_groups.remove(bpy.data.node_groups.get("Section Override"))
+    try:
-        bpy.data.node_groups.remove(bpy.data.node_groups.get("Section Compare"))
+        emission_node = next(n for n in section_node_group.nodes if isinstance(n, bpy.types.ShaderNodeEmission))
-        bpy.ops.object.delete({"selected_objects": [context.active_object]})
+        emission_node.inputs[0].default_value = list(self.section_plane_colour) + [1]
+    except StopIteration:
+        pass
 
 
 class SectionProperties(PropertyGroup):
     should_section_selected_objects: BoolProperty(name="Section Selected Objects", default=False)
+    hide: BoolProperty(default=False, name="Toggle", update=update_visibility)
     section_plane_colour: FloatVectorProperty(
         name="Temporary Section Cutaway Colour",
         subtype="COLOR",
@@ -304,7 +363,7 @@ class SectionProperties(PropertyGroup):
     )
 
 
-class BIM_PT_section_plane(bpy.types.Panel):
+class BIM_PT_section_plane(Panel):
     bl_label = "Temporary Section Cutaways"
     bl_idname = "BIM_PT_section_plane"
     bl_space_type = "VIEW_3D"
@@ -316,30 +375,31 @@ class BIM_PT_section_plane(bpy.types.Panel):
         layout.use_property_split = True
         props = context.scene.SectionProperties
 
-        row = layout.row()
+        layout.prop(props, "should_section_selected_objects", text="Only Selected", icon="RESTRICT_SELECT_OFF")
-        row.prop(props, "should_section_selected_objects")
+        layout.prop(props, "section_plane_colour", text="Section Color")
+        layout.prop(props, "hide", icon="HIDE_ON" if props.hide else "HIDE_OFF")
 
-        row = layout.row()
+        layout.operator("bim.section_plane_refresh", text="Refresh", icon="FILE_REFRESH")
-        row.prop(props, "section_plane_colour")
+        layout.operator("bim.section_plane_add", text="Add Section Cutaway", icon="ADD")
-
+        layout.operator("bim.section_plane_remove", text="Remove Selected Sections", icon="REMOVE").remove_all = False
-        row = layout.row(align=True)
+        layout.operator("bim.section_plane_remove", text="Remove All Sections", icon="TRASH").remove_all = True
-        row.operator("bim.add_section_plane")
-        row.operator("bim.remove_section_plane")
 
 
 def register():
     bpy.utils.register_class(SectionProperties)
-    bpy.utils.register_class(BIM_OT_add_section_plane)
+    bpy.utils.register_class(BIM_OT_section_plane_add)
-    bpy.utils.register_class(BIM_OT_remove_section_plane)
+    bpy.utils.register_class(BIM_OT_section_plane_remove)
     bpy.utils.register_class(BIM_PT_section_plane)
-    bpy.types.Scene.SectionProperties = bpy.props.PointerProperty(type=SectionProperties)
+    bpy.utils.register_class(BIM_OT_section_plane_refresh)
+    bpy.types.Scene.SectionProperties = PointerProperty(type=SectionProperties)
 
 
 def unregister():
-    bpy.utils.unregister_class(BIM_OT_add_section_plane)
+    bpy.utils.unregister_class(BIM_OT_section_plane_add)
-    bpy.utils.unregister_class(BIM_OT_remove_section_plane)
+    bpy.utils.unregister_class(BIM_OT_section_plane_remove)
     bpy.utils.unregister_class(BIM_PT_section_plane)
     bpy.utils.unregister_class(SectionProperties)
+    bpy.utils.unregister_class(BIM_OT_section_plane_refresh)
     del bpy.types.Scene.SectionProperties