diff --git a/compositor/operator.py b/compositor/operator.py index b1daada..8078912 100644 --- a/compositor/operator.py +++ b/compositor/operator.py @@ -2,17 +2,18 @@ from bpy.types import Node, CompositorNodeColorBalance, CompositorNodeTree -from ..ntp_operator import NTP_Operator +from ..ntp_operator import NTP_Operator, INDEX from ..ntp_node_tree import NTP_NodeTree from ..utils import * from io import StringIO from .node_settings import compositor_node_settings -SCENE_VAR = "scene" -BASE_NAME_VAR = "base_name" -END_NAME_VAR = "end_name" +SCENE = "scene" +BASE_NAME = "base_name" +END_NAME = "end_name" +NODE = "node" -ntp_vars = {SCENE_VAR, BASE_NAME_VAR, END_NAME_VAR} +comp_op_reserved_names = {SCENE, BASE_NAME, END_NAME, NODE} class NTPCompositorOperator(NTP_Operator): bl_idname = "node.ntp_compositor" @@ -33,31 +34,33 @@ class NTPCompositorOperator(NTP_Operator): def __init__(self): super().__init__() self._settings = compositor_node_settings + for name in comp_op_reserved_names: + self._used_vars[name] = 0 def _create_scene(self, indent: str): #TODO: wrap in more general unique name util function self._write(f"# Generate unique scene name", indent) - self._write(f"{BASE_NAME_VAR} = {str_to_py_str(self.compositor_name)}", + self._write(f"{BASE_NAME} = {str_to_py_str(self.compositor_name)}", indent) - self._write(f"{END_NAME_VAR} = {BASE_NAME_VAR}", indent) - self._write(f"if bpy.data.scenes.get({END_NAME_VAR}) != None:", indent) + self._write(f"{END_NAME} = {BASE_NAME}", indent) + self._write(f"if bpy.data.scenes.get({END_NAME}) != None:", indent) indent2 = f"{indent}\t" - self._write(f"i = 1", indent2) - self._write(f"{END_NAME_VAR} = {BASE_NAME_VAR} + f\".{{i:03d}}\"", + self._write(f"{INDEX} = 1", indent2) + self._write(f"{END_NAME} = {BASE_NAME} + f\".{{i:03d}}\"", indent2) - self._write(f"while bpy.data.scenes.get({END_NAME_VAR}) != None:", + self._write(f"while bpy.data.scenes.get({END_NAME}) != None:", indent2) indent3 = f"{indent}\t\t" - self._write(f"{END_NAME_VAR} = {BASE_NAME_VAR} + f\".{{i:03d}}\"", indent3) - self._write(f"i += 1\n", indent3) + self._write(f"{END_NAME} = {BASE_NAME} + f\".{{{INDEX}:03d}}\"", indent3) + self._write(f"{INDEX} += 1\n", indent3) - self._write(f"{SCENE_VAR} = bpy.context.window.scene.copy()\n", indent) - self._write(f"{SCENE_VAR}.name = {END_NAME_VAR}", indent) - self._write(f"{SCENE_VAR}.use_fake_user = True", indent) - self._write(f"bpy.context.window.scene = {SCENE_VAR}", indent) + self._write(f"{SCENE} = bpy.context.window.scene.copy()\n", indent) + self._write(f"{SCENE}.name = {END_NAME}", indent) + self._write(f"{SCENE}.use_fake_user = True", indent) + self._write(f"bpy.context.window.scene = {SCENE}", indent) def _initialize_compositor_node_tree(self, ntp_nt, nt_name): #initialize node group @@ -65,10 +68,10 @@ def _initialize_compositor_node_tree(self, ntp_nt, nt_name): self._write(f"def {ntp_nt.var}_node_group():", self._outer) if ntp_nt.node_tree == self._base_node_tree: - self._write(f"{ntp_nt.var} = {SCENE_VAR}.node_tree") + self._write(f"{ntp_nt.var} = {SCENE}.node_tree") self._write(f"#start with a clean node tree") - self._write(f"for node in {ntp_nt.var}.nodes:") - self._write(f"\t{ntp_nt.var}.nodes.remove(node)") + self._write(f"for {NODE} in {ntp_nt.var}.nodes:") + self._write(f"\t{ntp_nt.var}.nodes.remove({NODE})") else: self._write((f"{ntp_nt.var} = bpy.data.node_groups.new(" f"type = \'CompositorNodeTree\', " diff --git a/geometry/operator.py b/geometry/operator.py index 2f51c3a..9733a15 100644 --- a/geometry/operator.py +++ b/geometry/operator.py @@ -2,13 +2,6 @@ from bpy.types import GeometryNode, GeometryNodeTree from bpy.types import Node -if bpy.app.version >= (3, 6, 0): - from bpy.types import GeometryNodeSimulationInput - from bpy.types import GeometryNodeSimulationOutput -if bpy.app.version >= (4, 0, 0): - from bpy.types import GeometryNodeRepeatInput - from bpy.types import GeometryNodeRepeatOutput - from io import StringIO from ..ntp_operator import NTP_Operator @@ -16,6 +9,15 @@ from .node_tree import NTP_GeoNodeTree from .node_settings import geo_node_settings +ITEM = "item" +OBJECT_NAME = "name" +OBJECT = "obj" +MODIFIER = "mod" +geo_op_reserved_names = {ITEM, + OBJECT_NAME, + OBJECT, + MODIFIER} + class NTPGeoNodesOperator(NTP_Operator): bl_idname = "node.ntp_geo_nodes" bl_label = "Geo Nodes to Python" @@ -34,6 +36,8 @@ class NTPGeoNodesOperator(NTP_Operator): def __init__(self): super().__init__() self._settings = geo_node_settings + for name in geo_op_reserved_names: + self._used_vars[name] = 0 if bpy.app.version >= (3, 6, 0): def _process_zone_output_node(self, node: GeometryNode) -> None: @@ -50,15 +54,14 @@ def _process_zone_output_node(self, node: GeometryNode) -> None: node_var = self._node_vars[node] self._write(f"# Remove generated {items}") - self._write(f"for item in {node_var}.{items}:") + self._write(f"for {ITEM} in {node_var}.{items}:") self._write(f"\t{node_var}.{items}.remove(item)") for i, item in enumerate(getattr(node, items)): socket_type = enum_to_py_str(item.socket_type) name = str_to_py_str(item.name) self._write(f"# Create item {name}") - self._write(f"{node_var}.{items}.new" - f"({socket_type}, {name})") + self._write(f"{node_var}.{items}.new({socket_type}, {name})") if is_sim: item_var = f"{node_var}.{items}[{i}]" ad = enum_to_py_str(item.attribute_domain) @@ -199,14 +202,14 @@ def _process_node_tree(self, node_tree: GeometryNodeTree) -> None: def _apply_modifier(self, nt: GeometryNodeTree, nt_var: str): #get object - self._write(f"name = bpy.context.object.name", self._outer) - self._write(f"obj = bpy.data.objects[name]", self._outer) + self._write(f"{OBJECT_NAME} = bpy.context.object.name", self._outer) + self._write(f"{OBJECT} = bpy.data.objects[{OBJECT_NAME}]", self._outer) #set modifier to the one we just created mod_name = str_to_py_str(nt.name) - self._write(f"mod = obj.modifiers.new(name = {mod_name}, " + self._write(f"{MODIFIER} = obj.modifiers.new(name = {mod_name}, " f"type = 'NODES')", self._outer) - self._write(f"mod.node_group = {nt_var}", self._outer) + self._write(f"{MODIFIER}.node_group = {nt_var}", self._outer) def execute(self, context): diff --git a/material/operator.py b/material/operator.py index 97bc342..485b530 100644 --- a/material/operator.py +++ b/material/operator.py @@ -10,6 +10,8 @@ from .node_settings import shader_node_settings MAT_VAR = "mat" +NODE = "node" +shader_op_reserved_names = {MAT_VAR, NODE} class NTPMaterialOperator(NTP_Operator): bl_idname = "node.ntp_material" @@ -22,6 +24,8 @@ class NTPMaterialOperator(NTP_Operator): def __init__(self): super().__init__() self._settings = shader_node_settings + for name in shader_op_reserved_names: + self._used_vars[name] = 0 def _create_material(self, indent: str): self._write(f"{MAT_VAR} = bpy.data.materials.new(" @@ -44,8 +48,8 @@ def _initialize_shader_node_tree(self, ntp_node_tree: NTP_NodeTree, if ntp_node_tree.node_tree == self._base_node_tree: self._write(f"{ntp_node_tree.var} = {MAT_VAR}.node_tree") self._write(f"#start with a clean node tree") - self._write(f"for node in {ntp_node_tree.var}.nodes:") - self._write(f"\t{ntp_node_tree.var}.nodes.remove(node)") + self._write(f"for {NODE} in {ntp_node_tree.var}.nodes:") + self._write(f"\t{ntp_node_tree.var}.nodes.remove({NODE})") else: self._write((f"{ntp_node_tree.var} = bpy.data.node_groups.new(" f"type = \'ShaderNodeTree\', " diff --git a/ntp_operator.py b/ntp_operator.py index ab46e40..a7b1862 100644 --- a/ntp_operator.py +++ b/ntp_operator.py @@ -16,6 +16,17 @@ from .ntp_node_tree import NTP_NodeTree from .utils import * +INDEX = "i" +IMAGE_DIR_NAME = "imgs" +IMAGE_PATH = "image_path" +BASE_DIR = "base_dir" + +reserved_names = { + INDEX, + IMAGE_DIR_NAME, + IMAGE_PATH, + BASE_DIR + } class NTP_Operator(Operator): """ @@ -70,7 +81,6 @@ def __init__(self): self._outer: str = "" self._inner: str = "\t" - # Base node tree we're converting self._base_node_tree: NodeTree = None @@ -86,6 +96,9 @@ def __init__(self): # Dictionary used for setting node properties self._settings: dict[str, list[(str, ST)]] = {} + for name in reserved_names: + self._used_vars[name] = 0 + def _write(self, string: str, indent: str = None): if indent is None: indent = self._inner @@ -192,8 +205,6 @@ def _create_var(self, name: str) -> str: Parameters: name (str): basic string we'd like to create the variable name out of - used_vars (dict[str, int]): dictionary containing variable names and - usage counts Returns: clean_name (str): variable name for the node tree @@ -214,7 +225,7 @@ def _create_node(self, node: Node, node_tree_var: str) -> str: Initializes a new node with location, dimension, and label info Parameters: - node (bpy.types.Node): node to be copied + node (Node): node to be copied node_tree_var (str): variable name for the node tree Returns: node_var (str): variable name for the node @@ -251,7 +262,7 @@ def _set_settings_defaults(self, node: Node) -> None: Sets the defaults for any settings a node may have Parameters: - node (bpy.types.Node): the node object we're copying settings from + node (Node): the node object we're copying settings from node_var (str): name of the variable we're using for the node in our add-on """ if node.bl_idname not in self._settings: @@ -296,13 +307,11 @@ def _set_settings_defaults(self, node: Node) -> None: elif type == ST.MATERIAL: name = str_to_py_str(attr.name) self._write((f"if {name} in bpy.data.materials:")) - self._write((f"\t{node_var}.{attr_name} = " - f"bpy.data.materials[{name}]")) + self._write((f"\t{setting_str} = bpy.data.materials[{name}]")) elif type == ST.OBJECT: name = str_to_py_str(attr.name) self._write((f"if {name} in bpy.data.objects:")) - self._write((f"\t{node_var}.{attr_name} = " - f"bpy.data.objects[{name}]")) + self._write((f"\t{setting_str} = bpy.data.objects[{name}]")) elif type == ST.COLOR_RAMP: self._color_ramp_settings(node, attr_name) elif type == ST.CURVE_MAPPING: @@ -356,10 +365,9 @@ def _group_io_settings(self, node: Node, Set the settings for group input and output sockets Parameters: - node (bpy.types.Node) : group input/output node + node (Node) : group input/output node io (str): whether we're generating the input or output settings - node_tree_var (str): variable name of the generated node tree - node_tree (bpy.types.NodeTree): node tree that we're generating + ntp_node_tree (NTP_NodeTree): node tree that we're generating input and output settings for """ node_tree_var = ntp_node_tree.var @@ -715,7 +723,7 @@ def _set_output_defaults(self, node: Node) -> None: Some output sockets need default values set. It's rather annoying Parameters: - node (bpy.types.Node): node for the output we're setting + node (Node): node for the output we're setting """ # TODO: probably should define elsewhere output_default_nodes = {'ShaderNodeValue', @@ -772,8 +780,7 @@ def _color_ramp_settings(self, node: Node, color_ramp_name: str) -> None: color_ramp: bpy.types.ColorRamp = getattr(node, color_ramp_name) if not color_ramp: - raise ValueError( - f"No color ramp named \"{color_ramp_name}\" found") + raise ValueError(f"No color ramp named \"{color_ramp_name}\" found") node_var = self._node_vars[node] @@ -798,7 +805,7 @@ def _color_ramp_settings(self, node: Node, color_ramp_name: str) -> None: self._write((f"{ramp_str}.elements.remove" f"({ramp_str}.elements[0])")) for i, element in enumerate(color_ramp.elements): - element_var = f"{node_var}_cre_{i}" + element_var = self._create_var(f"{node_var}_cre_{i}") if i == 0: self._write(f"{element_var} = {ramp_str}.elements[{i}]") self._write(f"{element_var}.position = {element.position}") @@ -861,39 +868,64 @@ def _curve_mapping_settings(self, node: Node, # create curves for i, curve in enumerate(mapping.curves): - # TODO: curve function - self._write(f"#curve {i}") - curve_i = f"{node_var}_curve_{i}" - self._write(f"{curve_i} = " - f"{node_var}.{curve_mapping_name}.curves[{i}]") - - # Remove default points when CurveMap is initialized with more than - # two points (just CompositorNodeHueCorrect) - if (node.bl_idname == 'CompositorNodeHueCorrect'): - self._write((f"for i in range" - f"(len({curve_i}.points.values()) - 1, 1, -1):")) - self._write( - f"\t{curve_i}.points.remove({curve_i}.points[i])") - - for j, point in enumerate(curve.points): - # TODO: point function - point_j = f"{curve_i}_point_{j}" - - loc = point.location - loc_str = f"{loc[0]}, {loc[1]}" - if j < 2: - self._write(f"{point_j} = {curve_i}.points[{j}]") - self._write(f"{point_j}.location = ({loc_str})") - else: - self._write(f"{point_j} = {curve_i}.points.new({loc_str})") - - handle = enum_to_py_str(point.handle_type) - self._write(f"{point_j}.handle_type = {handle}") + self._create_curve_map(node, i, curve, curve_mapping_name) # update curve self._write(f"#update curve after changes") self._write(f"{mapping_var}.update()") + def _create_curve_map(self, node: Node, i: int, curve: bpy.types.CurveMap, + curve_mapping_name: str) -> None: + """ + Helper function to create the ith curve of a node's curve mapping + + Parameters: + node (Node): the node with a curve mapping + i (int): index of the CurveMap within the mapping + curve (bpy.types.CurveMap): the curve map to recreate + curve_mapping_name (str): attribute name of the recreated curve mapping + """ + node_var = self._node_vars[node] + + self._write(f"#curve {i}") + curve_i_var = self._create_var(f"{node_var}_curve_{i}") + self._write(f"{curve_i_var} = " + f"{node_var}.{curve_mapping_name}.curves[{i}]") + + # Remove default points when CurveMap is initialized with more than + # two points (just CompositorNodeHueCorrect) + if (node.bl_idname == 'CompositorNodeHueCorrect'): + self._write(f"for {INDEX} in range" + f"(len({curve_i_var}.points.values()) - 1, 1, -1):") + self._write(f"\t{curve_i_var}.points.remove(" + f"{curve_i_var}.points[{INDEX}])") + + for j, point in enumerate(curve.points): + self._create_curve_map_point(j, point, curve_i_var) + + def _create_curve_map_point(self, j: int, point: bpy.types.CurveMapPoint, + curve_i_var: str) -> None: + """ + Helper function to recreate a curve map point + + Parameters: + j (int): index of the point within the curve map + point (CurveMapPoint): point to recreate + curve_i_var (str): variable name of the point's curve map + """ + point_j_var = self._create_var(f"{curve_i_var}_point_{j}") + + loc = point.location + loc_str = f"{loc[0]}, {loc[1]}" + if j < 2: + self._write(f"{point_j_var} = {curve_i_var}.points[{j}]") + self._write(f"{point_j_var}.location = ({loc_str})") + else: + self._write(f"{point_j_var} = {curve_i_var}.points.new({loc_str})") + + handle = enum_to_py_str(point.handle_type) + self._write(f"{point_j_var}.handle_type = {handle}") + def _node_tree_settings(self, node: Node, attr_name: str) -> None: """ Processes node tree of group node if one is present @@ -951,13 +983,13 @@ def _load_image(self, img: bpy.types.Image, img_var: str) -> None: # TODO: convert to special variables self._write(f"#load image {img_str}") - self._write(f"base_dir = " + self._write(f"{BASE_DIR} = " f"os.path.dirname(os.path.abspath(__file__))") - self._write(f"image_path = " - f"os.path.join(base_dir, \"{IMAGE_DIR_NAME}\", " + self._write(f"{IMAGE_PATH} = " + f"os.path.join({BASE_DIR}, \"{IMAGE_DIR_NAME}\", " f"\"{img_str}\")") self._write(f"{img_var} = bpy.data.images.load" - f"(image_path, check_existing = True)") + f"({IMAGE_PATH}, check_existing = True)") # copy image settings self._write(f"#set image settings") diff --git a/utils.py b/utils.py index ec5d04d..5b3cb5d 100644 --- a/utils.py +++ b/utils.py @@ -7,8 +7,6 @@ import re from typing import Tuple -IMAGE_DIR_NAME = "imgs" - #node input sockets that are messy to set default values for dont_set_defaults = {'NodeSocketGeometry', 'NodeSocketShader',