Input socket types for nodes

[Surasia]

Description about the feature

When creating nodes, it doesn't seem possible to use the "default_value" attribute for a NodeSocket, as those are defined for subclasses of NodeSocket such as NodeSocketFloat.

For instance, this code accesses the second input of a math node:

subtract = create_node(self.node_tree.nodes, -400, 1083, ShaderNodeMath) # subtract is of type ShaderNodeMath
subtract.operation = "SUBTRACT" # this works
subtract.inputs[1].default_value = 1.0

This however throws a type error as .default_value may is not defined in the NodeSocket class. It's possible to work around this by overriding the type of the input:

subtract = create_node(self.node_tree.nodes, -400, 1083, ShaderNodeMath)
subtract.operation = "SUBTRACT"
_: NodeSocketFloat = subtract.inputs[1]
_.default_value = 1.0

Unfortunately, the documentation does not seem to specify the types for each input in a node socket, so I'm not sure if an implementation is possible.

Are you willing to contribute about this feature.

Yes, contributing to blender docs to specify input types.

[JonathanPlasse]

I would love having this feature.

I think you could do the following to implement this feature.

1. Find a way to get the inputs type information from the Blender Python console.
2. Update gen_external_modules_modfile.py to generate mod files with the inputs type information.
3. Find a way to encode the input type information in Python when you can use both integer or string indexing.

Something like this could work for 3.

class A:
    @overload
    def __getitem__(self, key: Literal[1]) -> Literal["a"]: ...
    @overload
    def __getitem__(self, key: Literal["a"]) -> Literal["a"]: ...
    @overload
    def __getitem__(self, key: Literal[2]) -> Literal["b"]: ...
    @overload
    def __getitem__(self, key: Literal["b"]) -> Literal["b"]: ...
    def __getitem__(self, key: int | str) -> Any: ...

What is describe above do not require a change to the Blender documentation.
There may be push back from the Blender project to do such a huge change to the documentation as it would require modifying the documentation of all nodes.
This could be done in a second time.

You may need to make a PR to Blender to surface the node type information

@nutti, may have more to say on this.

[Surasia]

Iterating over the node inputs does return the class of the socket at runtime, will try writing something that should be able to dump that type information into a definition file. Thanks for the tips!

[Surasia]

Incredibly embarrassed to have come up with this, but this seems to work?
One issue is that each node type is not initialized when blender launches, so you need to go through the UI to register each. I'm not familiar with how blender's RNA system works so I'm sure there's a much better way.

import bpy

for node in bpy.types.ShaderNode.__subclasses__():
    obj = bpy.context.selected_objects[0].material_slots[0].material.node_tree
    m = obj.nodes.new(node.__name__)
    print("--------------")
    print(m.__class__.__name__)
    print("--------------")
    for i, input in enumerate(m.inputs):
        print(f"[{i}]: {input.__class__.__name__}: {input.name}")

[nutti]

@JonathanPlasse

I agree with your idea.
inputs whose data type is NodeInputs is common class used in many nodes.
https://docs.blender.org/api/current/bpy.types.Node.html#bpy.types.Node.inputs
So, generating mod file is a good solution to solve this issue.
But we need to concern about the overhead of applying mod files.

@Surasia
You can look into this issue.
Once you finished the work, please consider to make PR :)

[30350n]

Hi!

I've just ran into this and decided to look into fixing it. Based on the POC by @Surasia , I've written a simple generator script which adds a "private" _NodeInputs and _NodeOutputs class to each Node class and overloads the inputs: _NodeInputs and outputs: _NodeOutputs like so. The "private" types inherit from NodeInputs and NodeOutputs respectively and overload the __getitem__ methods as recommended by @JonathanPlasse .

Here's a small output example for ShaderNodeBsdfDiffuse (full output file):

class ShaderNodeBsdfDiffuse(ShaderNode, NodeInternal, Node, bpy_struct):
    # insert after the existing type hints

    class _NodeInputs(NodeInputs):
        @typing.overload
        def __getitem__(self, key: typing.Literal[0] | typing.Literal["Color"]) -> NodeSocketColor: ...
        @typing.overload
        def __getitem__(self, key: typing.Literal[1] | typing.Literal["Roughness"]) -> NodeSocketFloatFactor: ...
        @typing.overload
        def __getitem__(self, key: typing.Literal[2] | typing.Literal["Normal"]) -> NodeSocketVector: ...
        @typing.overload
        def __getitem__(self, key: typing.Literal[3] | typing.Literal["Weight"]) -> NodeSocketFloat: ...

    inputs: _NodeInputs

    class _NodeOutputs(NodeOutputs):
        @typing.overload
        def __getitem__(self, key: typing.Literal["BSDF"] | typing.Literal[0]) -> NodeSocketShader: ...

    outputs: _NodeOutputs

And here's the generator code:

from pathlib import Path

import bpy

OUTPUT_PATH = Path(__file__).parent / "node_typing_out.pyi"

with OUTPUT_PATH.open("w", encoding="utf-8") as output_file:
    for node_base_class in bpy.types.NodeInternal.__subclasses__():
        print(node_base_class)
        node_baseclass_name = node_base_class.__name__
        nodetree_id = f"{node_baseclass_name}Tree"
        nodetree_class = getattr(bpy.types, nodetree_id)
        node_group = bpy.data.node_groups.new(nodetree_id, nodetree_id)

        for node_class_name in dir(bpy.types):
            if not node_class_name.startswith(node_baseclass_name):
                continue

            try:
                node = node_group.nodes.new(node_class_name)
            except RuntimeError:
                continue

            output_file.write(
                f"class {node_class_name}({node_baseclass_name}, NodeInternal, Node, bpy_struct):\n"
                f"    # insert after the existing type hints\n\n"
            )

            if node.inputs:
                output_file.write("    class _NodeInputs(NodeInputs):\n")

                for i, input in enumerate(node.inputs):
                    output_file.write(
                        f"        @typing.overload\n"
                        f"        def __getitem__(self, key: typing.Literal[{i}]"
                        f' | typing.Literal["{input.name}"]) -> {input.__class__.__name__}: ...\n'
                    )

                output_file.write("\n    inputs: _NodeInputs\n\n")

            if node.outputs:
                output_file.write("    class _NodeOutputs(NodeOutputs):\n")

                for i, output in enumerate(node.outputs):
                    output_file.write(
                        f"        @typing.overload\n"
                        f'        def __getitem__(self, key: typing.Literal["{output.name}"]'
                        f" | typing.Literal[{i}]) -> {output.__class__.__name__}: ...\n"
                    )

                output_file.write("\n    outputs: _NodeOutputs\n\n")

@nutti I'd be happy to submit this as a pull request, but I'm not sure how/where this would be integrated into the existing generator code. Could you give me any pointers?

[nutti]

@30350n

Thank you for challenging this issue.

I think this should not output to mod files which requires lots of computing resources.
It is good to add transformer class for this purpose.
https://github.com/nutti/fake-bpy-module/tree/main/src/fake_bpy_module/transformer

Is it possible to create a custom transformer?

[30350n]

I've read through Generate Modules now to try to understand how this project functions and I'm not sure about that.
From what I understand, the bulk of the type hints are being generated directly from the sphinx documentation using said transformers, right?

The script I came up with requires information which is not present in the docs, so it has to be run from within Blender. Looking at gen_external_modules_modfile.py, this looks more similar to what I'd have to do here, do you agree?

Also, out of curiosity, what's the performance problem with "mod files"?

[nutti]

@30350n

OK, I misunderstood about it.
In the case, mod file option is the only way.
I think creating a new script is a better way to do it because this script does not target to the external module.