-
Notifications
You must be signed in to change notification settings - Fork 36
Expand file tree
/
Copy pathnode_group_gatherer.py
More file actions
134 lines (114 loc) · 4.52 KB
/
Copy pathnode_group_gatherer.py
File metadata and controls
134 lines (114 loc) · 4.52 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
import bpy
from enum import Enum, auto
class NodeGroupType(Enum):
COMPOSITOR_NODE_GROUP = auto()
SCENE = auto()
GEOMETRY_NODE_GROUP = auto()
LIGHT = auto()
LINE_STYLE = auto()
MATERIAL = auto()
SHADER_NODE_GROUP = auto()
WORLD = auto()
def is_group(self) -> bool:
return self in {
NodeGroupType.COMPOSITOR_NODE_GROUP,
NodeGroupType.GEOMETRY_NODE_GROUP,
NodeGroupType.SHADER_NODE_GROUP
}
def is_obj(self) -> bool:
return (not self.is_group())
def is_compositor(self) -> bool:
return self in {
NodeGroupType.COMPOSITOR_NODE_GROUP,
NodeGroupType.SCENE
}
def is_geometry(self) -> bool:
return self in {
NodeGroupType.GEOMETRY_NODE_GROUP
}
def is_shader(self) -> bool:
return self in {
NodeGroupType.LIGHT,
NodeGroupType.LINE_STYLE,
NodeGroupType.MATERIAL,
NodeGroupType.SHADER_NODE_GROUP,
NodeGroupType.WORLD
}
NTPObject = (
bpy.types.NodeTree
| bpy.types.Scene
| bpy.types.Light
| bpy.types.FreestyleLineStyle
| bpy.types.Material
| bpy.types.World
)
def get_base_node_tree(
ntp_obj: NTPObject, group_type: NodeGroupType
) -> bpy.types.NodeTree:
if group_type.is_group():
return ntp_obj
elif group_type == NodeGroupType.SCENE and bpy.app.version >= (5, 0, 0):
return getattr(ntp_obj, "compositing_node_group")
else:
return getattr(ntp_obj, "node_tree")
class NodeGroupGatherer:
def __init__(self):
self.node_groups : dict[NodeGroupType, list[NTPObject]] = {
NodeGroupType.COMPOSITOR_NODE_GROUP : [],
NodeGroupType.SCENE : [],
NodeGroupType.GEOMETRY_NODE_GROUP : [],
NodeGroupType.LIGHT : [],
NodeGroupType.LINE_STYLE : [],
NodeGroupType.MATERIAL : [],
NodeGroupType.SHADER_NODE_GROUP : [],
NodeGroupType.WORLD : [],
}
def gather_node_groups(self, context: bpy.types.Context):
for group_slot in getattr(context.scene, "ntp_compositor_node_group_slots"):
if group_slot.node_tree is not None:
self.node_groups[NodeGroupType.COMPOSITOR_NODE_GROUP].append(
group_slot.node_tree
)
for scene_slot in getattr(context.scene, "ntp_scene_slots"):
if scene_slot.scene is not None:
self.node_groups[NodeGroupType.SCENE].append(scene_slot.scene)
for group_slot in getattr(context.scene, "ntp_geometry_node_group_slots"):
if group_slot.node_tree is not None:
self.node_groups[NodeGroupType.GEOMETRY_NODE_GROUP].append(
group_slot.node_tree
)
for light_slot in getattr(context.scene, "ntp_light_slots"):
if light_slot.light is not None:
self.node_groups[NodeGroupType.LIGHT].append(light_slot.light)
for line_style_slot in getattr(context.scene, "ntp_line_style_slots"):
if line_style_slot.line_style is not None:
self.node_groups[NodeGroupType.LINE_STYLE].append(
line_style_slot.line_style
)
for material_slot in getattr(context.scene, "ntp_material_slots"):
if material_slot.material is not None:
self.node_groups[NodeGroupType.MATERIAL].append(
material_slot.material
)
for group_slot in getattr(context.scene, "ntp_shader_node_group_slots"):
if group_slot.node_tree is not None:
self.node_groups[NodeGroupType.SHADER_NODE_GROUP].append(
group_slot.node_tree
)
for world_slot in getattr(context.scene, "ntp_world_slots"):
if world_slot.world is not None:
self.node_groups[NodeGroupType.WORLD].append(world_slot.world)
def get_number_node_groups(self) -> int:
result = 0
for lst in self.node_groups.values():
result += len(lst)
return result
def get_single_node_group(self):
# TODO: better name/ergonomics in general
if self.get_number_node_groups() != 1:
raise ValueError("Expected just one node group")
for group_type, lst in self.node_groups.items():
if len(lst) == 1:
return group_type, lst[0]
raise AssertionError("Expected this to be unreachable")
classes = []