Add enum generation

api/node_mapper.py

@@ -1,5 +1,7 @@
 import bpy
 import bl_ui
+import itertools
+import enum
 from .state import State
 from .types import *
 from ..absolute_path import absolute_path
@@ -17,21 +19,39 @@ def build_node(node_type):
         for prop in node.bl_rna.properties:
             argname = prop.identifier.lower().replace(' ', '_')
             if argname in kwargs:
-                setattr(node, prop.identifier, kwargs[argname])
+                value = kwargs[argname]
+                if isinstance(value, enum.Enum):
+                    value = value.value
+                setattr(node, prop.identifier, value)
         for node_input in (node.inputs[1:] if _primary_arg is not None else node.inputs):
             argname = node_input.name.lower().replace(' ', '_')
+            all_with_name = []
+            for node_input2 in (node.inputs[1:] if _primary_arg is not None else node.inputs):
+                if node_input2.name.lower().replace(' ', '_') == argname and node_input2.type == node_input.type:
+                    all_with_name.append(node_input2)
             if argname in kwargs:
-                if node_input.is_multi_input and hasattr(kwargs[argname], '__iter__') and len(kwargs[argname]) > 0 and issubclass(type(next(iter(kwargs[argname]))), Type):
-                    for x in kwargs[argname]:
+                def set_or_create_link(x, node_input):
+                    if issubclass(type(x), Type):
                         State.current_node_tree.links.new(x._socket, node_input)
-                elif issubclass(type(kwargs[argname]), Type):
-                    State.current_node_tree.links.new(kwargs[argname]._socket, node_input)
+                    else:
+                        try:
+                            node_input.default_value = x
+                        except:
+                            constant = Type(value=x)
+                            State.current_node_tree.links.new(constant._socket, node_input)
+                value = kwargs[argname]
+                if isinstance(value, enum.Enum):
+                    value = value.value
+                if node_input.is_multi_input and hasattr(value, '__iter__') and len() > 0 and issubclass(type(next(iter(value))), Type):
+                    for x in value:
+                        for node_input in all_with_name:
+                            State.current_node_tree.links.new(x._socket, node_input)
+                elif len(all_with_name) > 1 and issubclass(type(value), tuple) and len(value) > 0:
+                    for i, x in enumerate(value):
+                        set_or_create_link(x, all_with_name[i])
                 else:
-                    try:
-                        node_input.default_value = kwargs[argname]
-                    except:
-                        constant = Type(value=kwargs[argname])
-                        State.current_node_tree.links.new(constant._socket, node_input)
+                    for node_input in all_with_name:
+                        set_or_create_link(value, node_input)
         outputs = {}
         for node_output in node.outputs:
             if not node_output.enabled:
@@ -49,6 +69,7 @@ def register_node(node_type, category_path=None):
     if node_type in registered_nodes:
         return
     snake_case_name = node_type.bl_rna.name.lower().replace(' ', '_')
+    node_namespace_name = snake_case_name.replace('_', ' ').title().replace(' ', '')
     globals()[snake_case_name] = build_node(node_type)
     globals()[snake_case_name].bl_category_path = category_path
     globals()[snake_case_name].bl_node_type = node_type
@@ -58,6 +79,16 @@ def register_node(node_type, category_path=None):
             return build_node(node_type)(self, *args, **kwargs)
         return build
     setattr(Type, snake_case_name, build_node_method(node_type))
+    parent_props = [prop.identifier for base in node_type.__bases__ for prop in base.bl_rna.properties]
+    for prop in node_type.bl_rna.properties:
+        if not prop.identifier in parent_props and prop.type == 'ENUM':
+            if node_namespace_name not in globals():
+                class NodeNamespace: pass
+                NodeNamespace.__name__ = node_namespace_name
+                globals()[node_namespace_name] = NodeNamespace
+            enum_type_name = prop.identifier.replace('_', ' ').title().replace(' ', '')
+            enum_type = enum.Enum(enum_type_name, { map_case_name(i): i.identifier for i in prop.enum_items })
+            setattr(globals()[node_namespace_name], enum_type_name, enum_type)
     registered_nodes.add(node_type)
 for category_name in list(filter(lambda x: x.startswith('NODE_MT_category_GEO_'), dir(bpy.types))):
     category = getattr(bpy.types, category_name)
@@ -108,6 +139,7 @@ def create_documentation():
     default_color = '#A1A1A1'
     docstrings = []
     symbols = []
+    enums = {}
     for func in sorted(documentation.keys()):
         try:
             method = documentation[func]
@@ -117,15 +149,21 @@ def create_documentation():
             props_inputs = {}
             symbol_inputs = {}
             parent_props = [prop.identifier for base in method.bl_node_type.__bases__ for prop in base.bl_rna.properties]
+            node_namespace_name = func.replace('_', ' ').title().replace(' ', '')
             for prop in method.bl_node_type.bl_rna.properties:
                 if not prop.identifier in parent_props:
                     if prop.type == 'ENUM':
-                        enum_items = 'Literal[' + ', '.join(map(lambda i: f"'{i.identifier}'", prop.enum_items)) + ']'
-                        props_inputs[prop.identifier] = f"<span style=\"color: {color_mappings['STRING']};\">{enum_items}</span>"
-                        symbol_inputs[prop.identifier] = enum_items
+                        enum_name = prop.identifier.replace('_', ' ').title().replace(' ', '')
+                        enum_cases = '\n    '.join(map(lambda i: f"{map_case_name(i)} = '{i.identifier}'", prop.enum_items))
+                        if node_namespace_name not in enums:
+                            enums[node_namespace_name] = []
+                        enums[node_namespace_name].append(f"""  class {enum_name}(enum.Enum):
+    {enum_cases}""")
+                        props_inputs[prop.identifier] = {f"<span style=\"color: {color_mappings['STRING']};\">{node_namespace_name}.{enum_name}</span>":1}
+                        symbol_inputs[prop.identifier] = {f"{node_namespace_name}.{enum_name}": 1}
                     else:
-                        props_inputs[prop.identifier] = f"<span style=\"color: {color_mappings.get(prop.type, default_color)};\">{prop.type.title()}</span>"
-                        symbol_inputs[prop.identifier] = prop.type.title()
+                        props_inputs[prop.identifier] = {f"<span style=\"color: {color_mappings.get(prop.type, default_color)};\">{prop.type.title()}</span>":1}
+                        symbol_inputs[prop.identifier] = {prop.type.title(): 1}
             primary_arg = None
             for node_input in node_instance.inputs:
                 name = node_input.name.lower().replace(' ', '_')
@@ -134,13 +172,32 @@ def create_documentation():
                     typename = f"List[{typename}]"
                 type_str = f"<span style=\"color: {color_mappings.get(node_input.type, default_color)};\">{typename}</span>"
                 if name in props_inputs:
-                    props_inputs[name] = props_inputs[name] + f' | {type_str}'
-                    symbol_inputs[name] = symbol_inputs[name] + f' | {typename}'
+                    if type_str in props_inputs[name]:
+                        props_inputs[name][type_str] += 1
+                        symbol_inputs[name][typename] += 1
+                    else:
+                        props_inputs[name][type_str] = 1
+                        symbol_inputs[name][typename] = 1
                 else:
-                    props_inputs[name] = type_str
-                    symbol_inputs[name] = typename
+                    props_inputs[name] = {type_str: 1}
+                    symbol_inputs[name] = {typename: 1}
                 if primary_arg is None:
-                    primary_arg = (name, props_inputs[name])
+                    primary_arg = (name, list(props_inputs[name].keys())[0])
+            def collapse_inputs(inputs):
+                for k, v in inputs.items():
+                    values = []
+                    for t, c in v.items():
+                        for c in range(1, c + 1):
+                            value = ""
+                            if c > 1:
+                                value += "Tuple["
+                            value += ', '.join(itertools.repeat(t, c))
+                            if c > 1:
+                                value += "]"
+                            values.append(value)
+                    inputs[k] = ' | '.join(values)
+            collapse_inputs(props_inputs)
+            collapse_inputs(symbol_inputs)
             arg_docs = []
             symbol_args = []
             for name, value in props_inputs.items():
@@ -181,14 +238,20 @@ def create_documentation():
                 </div>
             </details>
             """)
-            output_symbol_separator = '\n  '
-            symbol_return_type = f"_{func}_result"
+            output_symbol_separator = '\n    '
             if len(output_symbols) > 1:
-                symbols.append(f"""class {symbol_return_type}:
-  {output_symbol_separator.join(output_symbols)}""")
-            return_type_hint = list(symbol_outputs.values())[0] if len(output_symbols) == 1 else symbol_return_type
+                if node_namespace_name not in enums:
+                    enums[node_namespace_name] = []
+                enums[node_namespace_name].append(f"""  class Result:
+    {output_symbol_separator.join(output_symbols)}""")
+            return_type_hint = list(symbol_outputs.values())[0] if len(output_symbols) == 1 else f"{node_namespace_name}.Result"
             symbols.append(f"""def {func}({', '.join(symbol_args)}) -> {return_type_hint}: \"\"\"![]({image}.webp)\"\"\"""")
-        except:
+        except Exception as e:
+            import os, sys
+            print(e)
+            exc_type, exc_obj, exc_tb = sys.exc_info()
+            fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
+            print(exc_type, fname, exc_tb.tb_lineno)
             continue
     bpy.data.node_groups.remove(temp_node_group)
     html = f"""
@@ -228,15 +291,40 @@ def create_documentation():
         newline = '\n'
         def type_symbol(t):
             return f"class {t.__name__}(Type): pass"
+        def enum_namespace(k):
+            return f"""class {k}:
+{newline.join(enums[k])}"""
         contents = f"""from typing import *
+import enum
 def tree(builder):
   \"\"\"
   Marks a function as a node tree.
   \"\"\"
   pass
 class Type:
-  {(newline + '  ').join(filter(lambda x: x.startswith('def'), symbols))}
+  def __add__(self, other) -> Type: return self
+  def __radd__(self, other) -> Type: return self
+  def __sub__(self, other) -> Type: return self
+  def __rsub__(self, other) -> Type: return self
+  def __mul__(self, other) -> Type: return self
+  def __rmul__(self, other) -> Type: return self
+  def __truediv__(self, other) -> Type: return self
+  def __rtruediv__(self, other) -> Type: return self
+  def __mod__(self, other) -> Type: return self
+  def __rmod__(self, other) -> Type: return self
+  def __eq__(self, other) -> Type: return self
+  def __ne__(self, other) -> Type: return self
+  def __lt__(self, other) -> Type: return self
+  def __le__(self, other) -> Type: return self
+  def __gt__(self, other) -> Type: return self
+  def __ge__(self, other) -> Type: return self
+  x = Type()
+  y = Type()
+  z = Type()
+  {(newline + '  ').join(map(lambda x: x.replace('(', '(self, '), filter(lambda x: x.startswith('def'), symbols)))}
+  
 {newline.join(map(type_symbol, Type.__subclasses__()))}
+{newline.join(map(enum_namespace, enums.keys()))}
 {newline.join(symbols)}"""
         fpyi.write(contents)
         fpy.write(contents)

api/types.py

@@ -3,6 +3,15 @@ from bpy.types import NodeSocketStandard
 import nodeitems_utils
 from .state import State
 
+def map_case_name(i):
+    r = i.identifier.replace('_', ' ').title().replace(' ', '')
+    if r == 'None':
+        return 'NONE'
+    elif not r[0].isalpha():
+        return f'_{r}'
+    else:
+        return r
+
 # The base class all exposed socket types conform to.
 class Type:
     socket_type: str
@@ -28,31 +37,46 @@ class Type:
         self._socket = socket
         self.socket_type = type(socket).__name__
     
-    def _math(self, other, operation):
+    def _math(self, other, operation, reverse=False):
         math_node = State.current_node_tree.nodes.new('ShaderNodeVectorMath' if self._socket.type == 'VECTOR' else 'ShaderNodeMath')
         math_node.operation = operation
-        State.current_node_tree.links.new(self._socket, math_node.inputs[0])
+        State.current_node_tree.links.new(self._socket, math_node.inputs[1 if reverse else 0])
         if issubclass(type(other), Type):
-            State.current_node_tree.links.new(other._socket, math_node.inputs[1])
+            State.current_node_tree.links.new(other._socket, math_node.inputs[0 if reverse else 1])
         else:
-            math_node.inputs[1].default_value = other
+            math_node.inputs[0 if reverse else 1].default_value = other
         return Type(math_node.outputs[0])
 
     def __add__(self, other):
         return self._math(other, 'ADD')
     
+    def __radd__(self, other):
+        return self._math(other, 'ADD', True)
+    
     def __sub__(self, other):
         return self._math(other, 'SUBTRACT')
     
+    def __rsub__(self, other):
+        return self._math(other, 'SUBTRACT', True)
+    
     def __mul__(self, other):
         return self._math(other, 'MULTIPLY')
     
+    def __rmul__(self, other):
+        return self._math(other, 'MULTIPLY', True)
+    
     def __truediv__(self, other):
         return self._math(other, 'DIVIDE')
     
+    def __rtruediv__(self, other):
+        return self._math(other, 'DIVIDE', True)
+    
     def __mod__(self, other):
         return self._math(other, 'MODULO')
     
+    def __rmod__(self, other):
+        return self._math(other, 'MODULO', True)
+    
     def _compare(self, other, operation):
         compare_node = State.current_node_tree.nodes.new('FunctionNodeCompare')
         compare_node.data_type = 'FLOAT' if self._socket.type == 'VALUE' else self._socket.type

book/book.toml

@@ -4,3 +4,8 @@ language = "en"
 multilingual = false
 src = "src"
 title = "Geometry Script"
+
+[output.html]
+default-theme = "coal"
+preferred-dark-theme = "coal"
+additional-css = ["style.css"]

book/src/api/basics/sockets.md

@@ -121,4 +121,30 @@ The same script without chaining calls is written more verbosely as:
 @tree("Cube Tree")
 def cube_tree(size: Vector):
     return mesh_to_volume(mesh=cube(size=size))
+```
+
+### Spanning Multiple Lines
+
+Often times you want each chained calls to be on a separate line. There are a few ways to do this in Python:
+
+1. Newlines around arguments
+
+```python
+cube(
+    size=size
+).mesh_to_volume()
+```
+
+2. Parentheses
+
+```python
+(cube(size=size)
+ .mesh_to_volume())
+```
+
+3. Line continuation
+
+```python
+cube(size=size) \
+    .mesh_to_volume()
 ```

book/src/api/basics/using-nodes.md

@@ -20,7 +20,47 @@ The general process is:
 
 > Properties and inputs are different types of argument. A property is a value that cannot be connected to a socket. These are typically enums (displayed in the UI as a dropdown), with specific string values expected. Check the documentation for a node to see what the possible values are for a property.
 
-Let's take a look at two nodes as an example.
+## Enum Properties
+
+Many nodes have enum properties. For example, the math node lets you choose which operation to perform. You can pass a string to specify the enum case to use. But a safer way to set these values is with the autogenerated enum types. The enums are namespaced to the name of the node in PascalCase:
+
+```python
+# Access it by Node.Enum Name.Case
+math(operation=Math.Operation.Add)
+math(operation=Math.Operation.Subtract)
+math(operation='MULTIPLY') # Or manually pass a string
+```
+
+Internally, this type is generated as:
+
+```python
+import enum
+class Math:
+  class Operation(enum.Enum):
+    Add = 'ADD'
+    Subtract = 'SUBTRACT'
+    Multiply = 'MULTIPLY'
+    ...
+  ...
+```
+
+The cases will appear in code completion if you setup an [external editor](../../setup/external-editing.md).
+
+## Duplicate Names
+
+Some nodes use the same input name multiple times. For example, the *Math* node has three inputs named `value`. To specify each value, pass a tuple for the input:
+
+```python
+math(operation=Math.Operation.Wrap, value=(0.5, 1, 0)) # Pass all 3
+math(operation=Math.Operation.Wrap, value=(0.5, 1)) # Only pass 2/3
+math(operation=Math.Operation.Wrap, value=0.5) # Only pass 1/3
+```
+
+![](./math_wrap.png)
+
+## Examples
+
+Here are two examples to show how a node maps to a function.
 
 ### Cube
 
@@ -72,8 +112,8 @@ size.cube(...)
 1. Name `Capture Attribute` -> `capture_attribute`
 2. Keyword Arguments
     * Properties
-        * `data_type: Literal['FLOAT', 'INT', 'FLOAT_VECTOR', 'FLOAT_COLOR', 'BYTE_COLOR', 'STRING', 'BOOLEAN', 'FLOAT2', 'INT8']`
-        * `domain: Literal['POINT', 'EDGE', 'FACE', 'CORNER', 'CURVE', 'INSTANCE']`
+        * `data_type: CaptureAttribute.DataType`
+        * `domain: CaptureAttribute.Domain`
     * Inputs
         * `geometry: Geometry`
         * `value: Vector | Float | Color | Bool | Int`
@@ -82,7 +122,7 @@ size.cube(...)
 The node can now be used as a function:
 
 ```python
-result = capture_attribute(data_type='BOOLEAN', geometry=cube_geo) # Specify a property and an input
+result = capture_attribute(data_type=CaptureAttribute.DataType.Boolean, geometry=cube_geo) # Specify a property and an input
 result.geometry # Access the geometry
 result.attribute # Access the attribute
 ```
@@ -92,8 +132,8 @@ The generated documentation will show the signature, result type, and [chain syn
 #### Signature
 ```python
 capture_attribute(
-  data_type: Literal['FLOAT', 'INT', 'FLOAT_VECTOR', 'FLOAT_COLOR', 'BYTE_COLOR', 'STRING', 'BOOLEAN', 'FLOAT2', 'INT8'],
-  domain: Literal['POINT', 'EDGE', 'FACE', 'CORNER', 'CURVE', 'INSTANCE'],
+  data_type: CaptureAttribute.DataType,
+  domain: CaptureAttribute.Domain,
   geometry: Geometry,
   value: Vector | Float | Color | Bool | Int
 )

book/src/setup/external-editing.md

@@ -3,9 +3,9 @@
 Blender's *Text Editor* leaves a lot to be desired. Writing scripts without code completion can be tough.
 Using an external code editor is one way to improve the editing experience.
 
-This guide will show how to setup [Visual Studio Code](https://code.visualstudio.com/) to edit Geometry Scripts. However, the same concepts apply to IDEs.
+This guide will show how to setup [Visual Studio Code](https://code.visualstudio.com/) to edit Geometry Scripts. However, the same concepts apply to other IDEs.
 
-> This guide assumes you have already installed Visual Studio Code and setup the [Python extension](https://marketplace.visualstudio.com/items?itemName=ms-python.python). If not, please follow the setup guides for those tools before continuing.
+> This guide assumes you have already installed Visual Studio Code and setup the [Python extension](https://marketplace.visualstudio.com/items?itemName=ms-python.python). If not, please setup those tools before continuing.
 
 ## Code Completion
 When the Geometry Script add-on starts, it generates a Python typeshed file that can be used to provide code completion.

book/src/tutorials/city-builder.md

@@ -26,7 +26,7 @@ def city_builder(
     return geometry
 ```
 
-Run the script to create the tree, then add a *Geometry Nodes* modifier to your curve object and select the *City Builger* node group.
+Run the script to create the tree, then add a *Geometry Nodes* modifier to your curve object and select the *City Builder* node group.
 
 ## Buildings
 Let's start with the buildings. We'll distribute points on a grid with `size_x` and `size_y`.

book/style.css

@@ -0,0 +1,3 @@
+.coal {
+    --bg: #1C1C1C !important;
+}