Merge pull request #135 from BrendanParmer/MenuSocketFixes

Menu socket fixes

Author: BrendanParmer

CONTRIBUTORS.md

@@ -0,0 +1,3 @@
+* [Brendan Parmer](https://github.com/BrendanParmer)
+* [Carlsu](https://github.com/carls3d)
+* [atticus-lv](https://github.com/atticus-lv)

NodeToPython/ntp_operator.py

@@ -12,7 +12,7 @@
 import datetime
 import os
 import shutil
-from typing import TextIO
+from typing import TextIO, Callable
 
 from .license_templates import license_templates
 from .ntp_node_tree import NTP_NodeTree
@@ -66,6 +66,9 @@ class NTP_Operator(Operator):
     def __init__(self):
         super().__init__()
 
+        # Write functions after nodes are mostly initialized and linked up
+        self._write_after_links: list[Callable] = []
+
         # File (TextIO) or string (StringIO) the add-on/script is generated into
         self._file: TextIO = None
 
@@ -538,7 +541,21 @@ def _set_tree_socket_defaults(self, socket_interface: NodeTreeInterfaceSocket,
 
             dv = socket_interface.default_value
 
-            if type(socket_interface) == bpy.types.NodeTreeInterfaceSocketColor:
+            if type(socket_interface) is bpy.types.NodeTreeInterfaceSocketMenu:
+                if dv == "":
+                    self.report({'WARNING'},
+                        "NodeToPython: No menu found for socket "
+                        f"{socket_interface.name}"
+                    )
+                    return
+
+                self._write_after_links.append(
+                    lambda _socket_var=socket_var, _dv=enum_to_py_str(dv): (
+                        self._write(f"{_socket_var}.default_value = {_dv}")
+                    )
+                )
+                return
+            elif type(socket_interface) == bpy.types.NodeTreeInterfaceSocketColor:
                 dv = vec4_to_py_str(dv)
             elif type(dv) in {mathutils.Vector, mathutils.Euler}:
                 dv = vec3_to_py_str(dv)
@@ -784,6 +801,8 @@ def _set_input_defaults(self, node: Node) -> None:
 
                 #menu
                 elif input.bl_idname == 'NodeSocketMenu':
+                    if input.default_value == '':
+                        continue
                     default_val = enum_to_py_str(input.default_value)
 
                 # images
@@ -1329,6 +1348,11 @@ def _init_links(self, node_tree: NodeTree) -> None:
                         f".outputs[{input_idx}], "
                         f"{out_node_var}.inputs[{output_idx}])")
 
+        for _func in self._write_after_links:
+            _func()
+        self._write_after_links = []
+            
+
     def _set_node_tree_properties(self, node_tree: NodeTree) -> None:
         nt_var = self._node_tree_vars[node_tree]