Skip to content

Commit 1dda44e

Browse files
BrendanParmerBrendanParmer
committed
feat: image saving and zipping
1 parent a8ff077 commit 1dda44e

3 files changed

Lines changed: 106 additions & 16 deletions

File tree

‎geo_nodes.py‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
#node input sockets that are messy to set default values for
1414
dont_set_defaults = {'NodeSocketCollection',
1515
'NodeSocketGeometry',
16-
'NodeSocketImage',
16+
#'NodeSocketImage',
1717
'NodeSocketMaterial',
1818
'NodeSocketObject',
1919
'NodeSocketTexture',
@@ -153,18 +153,18 @@ def execute(self, context):
153153
class_name = nt.name.replace(" ", "").replace('.', "")
154154

155155
#find base directory to save new addon
156-
dir = bpy.path.abspath("//")
157-
if not dir or dir == "":
156+
base_dir = bpy.path.abspath("//")
157+
if not base_dir or base_dir == "":
158158
self.report({'ERROR'},
159159
("NodeToPython: Save your blend file before using "
160160
"NodeToPython!"))
161161
return {'CANCELLED'}
162162

163163
#save in /addons/ subdirectory
164-
addon_dir = os.path.join(dir, "addons")
164+
addon_dir = os.path.join(base_dir, "addons", nt_var)
165165
if not os.path.exists(addon_dir):
166166
os.mkdir(addon_dir)
167-
file = open(f"{addon_dir}/{nt_var}_addon.py", "w")
167+
file = open(f"{addon_dir}/__init__.py", "w")
168168

169169
create_header(file, nt)
170170
init_operator(file, class_name, nt_var, nt.name)
@@ -318,7 +318,7 @@ def process_geo_nodes_group(node_tree, level):
318318
curve_node_settings(node, file, inner, node_var)
319319

320320
set_input_defaults(node, dont_set_defaults, file, inner,
321-
node_var)
321+
node_var, addon_dir)
322322

323323
set_parents(node_tree, file, inner, node_vars)
324324
set_locations(node_tree, file, inner, node_vars)
@@ -339,6 +339,9 @@ def process_geo_nodes_group(node_tree, level):
339339
create_main_func(file)
340340

341341
file.close()
342+
343+
zip_addon(addon_dir)
344+
342345
return {'FINISHED'}
343346

344347
class SelectGeoNodesMenu(bpy.types.Menu):

‎materials.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
#node input sockets that are messy to set default values for
77
dont_set_defaults = {'NodeSocketCollection',
88
'NodeSocketGeometry',
9-
'NodeSocketImage',
9+
#'NodeSocketImage',
1010
'NodeSocketMaterial',
1111
'NodeSocketObject',
1212
'NodeSocketShader',
@@ -108,10 +108,10 @@ def execute(self, context):
108108
("NodeToPython: Save your blender file before using "
109109
"NodeToPython!"))
110110
return {'CANCELLED'}
111-
addon_dir = os.path.join(dir, "addons")
111+
addon_dir = os.path.join(dir, "addons", mat_var)
112112
if not os.path.exists(addon_dir):
113113
os.mkdir(addon_dir)
114-
file = open(f"{addon_dir}/{mat_var}_addon.py", "w")
114+
file = open(f"{addon_dir}/__init__.py", "w")
115115

116116
create_header(file, nt)
117117
init_operator(file, class_name, mat_var, self.material_name)
@@ -177,7 +177,7 @@ def process_mat_node_group(node_tree, level):
177177
curve_node_settings(node, file, inner, node_var)
178178

179179
set_input_defaults(node, dont_set_defaults, file, inner,
180-
node_var)
180+
node_var, addon_dir)
181181
set_parents(node_tree, file, inner, node_vars)
182182
set_locations(node_tree, file, inner, node_vars)
183183
set_dimensions(node_tree, file, inner, node_vars)
@@ -196,6 +196,7 @@ def process_mat_node_group(node_tree, level):
196196
create_main_func(file)
197197

198198
file.close()
199+
zip_addon(addon_dir)
199200
return {'FINISHED'}
200201

201202
class SelectMaterialMenu(bpy.types.Menu):

‎utils.py‎

Lines changed: 92 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
import bpy
22
import mathutils
33

4+
import os
45
import re
6+
import shutil
57
from typing import TextIO, Tuple
68

9+
image_dir_name = "imgs"
10+
711
def clean_string(string: str) -> str:
812
"""
913
Cleans up a string for use as a variable or file name
@@ -66,6 +70,20 @@ def vec4_to_py_str(vec) -> str:
6670
"""
6771
return f"({vec[0]}, {vec[1]}, {vec[2]}, {vec[3]})"
6872

73+
def img_to_py_str(img) -> str:
74+
"""
75+
Converts a Blender image into its string
76+
77+
Paramters:
78+
img (bpy.types.Image): a Blender image
79+
80+
Returns:
81+
(str): string version
82+
"""
83+
name = img.name
84+
format = img.file_format.lower()
85+
return f"{name}.{format}"
86+
6987
def create_header(file: TextIO, node_tree):
7088
"""
7189
Sets up the bl_info and imports the Blender API
@@ -85,6 +103,7 @@ def create_header(file: TextIO, node_tree):
85103
file.write("}\n")
86104
file.write("\n")
87105
file.write("import bpy\n")
106+
file.write("import os")
88107
file.write("\n")
89108

90109
def init_operator(file: TextIO, name: str, idname: str, label: str):
@@ -229,7 +248,6 @@ def color_ramp_settings(node, file: TextIO, inner: str, node_var: str):
229248
file.write((f"{inner}{node_var}.color_ramp.interpolation "
230249
f"= '{color_ramp.interpolation}'\n"))
231250
file.write("\n")
232-
233251
#key points
234252
for i, element in enumerate(color_ramp.elements):
235253
file.write((f"{inner}{node_var}_cre_{i} = "
@@ -305,22 +323,41 @@ def curve_node_settings(node, file: TextIO, inner: str, node_var: str):
305323
file.write(f"{inner}#update curve after changes\n")
306324
file.write(f"{mapping}.update()\n")
307325

308-
def set_input_defaults(node, dont_set_defaults: dict, file: TextIO, inner: str,
309-
node_var: str):
326+
def set_input_defaults(node, dont_set_defaults: set, file: TextIO, inner: str,
327+
node_var: str, addon_dir: str):
328+
"""
329+
Sets defaults for input sockets
330+
331+
Parameters:
332+
node (bpy.types.Node): node we're setting inputs for
333+
dont_set_defaults (set): set of sockets we shouldn't attempt to set
334+
default values for
335+
file (TextIO): file we're generating the add-on into
336+
inner (str): indentation
337+
node_var (str): variable name we're using for the copied node
338+
addon_dir (str): directory of the add-on, for if we need to save other
339+
objects for the add-on
340+
"""
310341
if node.bl_idname != 'NodeReroute':
311342
for i, input in enumerate(node.inputs):
312343
if input.bl_idname not in dont_set_defaults:
344+
socket_var = f"{node_var}.inputs[{i}]"
313345
if input.bl_idname == 'NodeSocketColor':
314346
default_val = vec4_to_py_str(input.default_value)
315347
elif "Vector" in input.bl_idname:
316348
default_val = vec3_to_py_str(input.default_value)
317349
elif input.bl_idname == 'NodeSocketString':
318350
default_val = str_to_py_str(input.default_value)
351+
elif input.bl_idname == 'NodeSocketImage':
352+
img = input.default_value
353+
save_image(img, addon_dir)
354+
load_image(img, file, inner, f"{socket_var}.default_value")
355+
default_val = None
319356
else:
320357
default_val = input.default_value
321358
if default_val is not None:
322359
file.write(f"{inner}#{input.identifier}\n")
323-
file.write((f"{inner}{node_var}.inputs[{i}].default_value"
360+
file.write((f"{inner}{socket_var}.default_value"
324361
f" = {default_val}\n"))
325362
file.write("\n")
326363

@@ -372,7 +409,6 @@ def set_dimensions(node_tree, file: TextIO, inner: str, node_vars: dict):
372409
file.write((f"{inner}{node_var}.width, {node_var}.height "
373410
f"= {node.width}, {node.height}\n"))
374411

375-
376412
def init_links(node_tree, file: TextIO, inner: str, node_tree_var: str,
377413
node_vars: dict):
378414
"""
@@ -464,4 +500,54 @@ def create_main_func(file: TextIO):
464500
file (TextIO): file we're generating the add-on into
465501
"""
466502
file.write("if __name__ == \"__main__\":\n")
467-
file.write("\tregister()")
503+
file.write("\tregister()")
504+
505+
def save_image(img, addon_dir: str):
506+
"""
507+
Saves an image to an image directory of the add-on
508+
509+
Parameters:
510+
img (bpy.types.Image): image to be saved
511+
addon_dir (str): directory of the addon
512+
"""
513+
#create image dir if one doesn't exist
514+
img_dir = os.path.join(addon_dir, image_dir_name)
515+
if not os.path.exists(img_dir):
516+
os.mkdir(img_dir)
517+
518+
#save the image
519+
img_str = img_to_py_str(img)
520+
img_path = f"{img_dir}/{img_str}"
521+
if not os.path.exists(img_path):
522+
img.save_render(img_path)
523+
524+
def load_image(img, file: TextIO, inner: str, img_var: str):
525+
"""
526+
Loads an image from the add-on into a blend file and assigns it
527+
528+
Parameters:
529+
img (bpy.types.Image): Blender image from the original node group
530+
file (TextIO): file for the generated add-on
531+
inner (str): indentation string
532+
img_var (str): variable name to be used for the image
533+
"""
534+
img_str = img_to_py_str(img)
535+
536+
file.write(f"{inner}#load image {img_str}\n")
537+
file.write((f"{inner}base_dir = "
538+
f"os.path.dirname(os.path.abspath(__file__))\n"))
539+
file.write((f"{inner}image_path = "
540+
f"os.path.join(base_dir, \"{image_dir_name}\", "
541+
f"\"{img_str}\")\n"))
542+
file.write((f"{inner}{img_var} = "
543+
f"bpy.data.images.load(image_path, check_existing = True)\n"))
544+
545+
def zip_addon(addon_dir: str):
546+
"""
547+
Zips up the addon and removes the directory
548+
549+
Parameters:
550+
addon_dir (str): path to the directory of the addon
551+
"""
552+
shutil.make_archive(addon_dir, "zip", addon_dir)
553+
shutil.rmtree(addon_dir)

0 commit comments

Comments
 (0)