compositor node helper.

AddCompositorNode.py

import sys
import os
import re

# Following the instructions from https://wiki.blender.org/index.php/User:Sobotka/Minimal_Node_Code

def SourcePath(SourceFolder, path):
    return os.path.join(SourceFolder,path)

def ValidNode(source, node):
    print("Step 0: Determine if node already exists")
    sp = SourcePath(source,"source/blender/blenkernel/BKE_node.h")
    nodeline =  "#define CMP_NODE_{}".format(node.upper())
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    for x in range(0, lines):
        if nodeline in contents[x]:
            return False
    return True

def step1_updatebkenode(source, node):
    print("Step 1: Register the Unique Node ID Value in source/blender/blenkernel/BKE_node.h")
    sp = SourcePath(source,"source/blender/blenkernel/BKE_node.h")
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    for x in range(0, lines):
        if "#define CMP_NODE_" in contents[x]:
            index = x
    tokens = contents[index].split()
    lastid = tokens[len(tokens)-1]
    newid = int(lastid)+1
    newline =  "#define CMP_NODE_{}\t{}\n".format(node.upper(), newid)
    contents.insert(index+1, newline)
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step2_updatenodec(source, node):
    print("Step 2: Register the new node in source/blender/blenkernel/intern/node.c")
    sp = SourcePath(source,"source/blender/blenkernel/intern/node.c")
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    for x in range(0, lines):
        if "register_node_type_cmp_" in contents[x]:
            index = x
    newline =  "\tregister_node_type_cmp_{}();\n".format(node.lower())
    contents.insert(index+1, newline)
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step3_updatedna(source, node):
    print("Step 3: Add the DNA that describes the data in source/blender/makesdna/DNA_node_types.h")
    sp = SourcePath(source,"source/blender/makesdna/DNA_node_types.h")
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    for x in range(0, lines):
        if "/* script node mode */" in contents[x]:
            index = x

    newline =  "typedef struct Node{} {{\n\tdouble {}Double; /* Our {} data. */\n}} Node{};\n\n".format(node,node,node,node)
    contents.insert(index, newline)
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step4_registerrna(source, node):
    print("Step 4: Register the Python RNA in source/blender/makesrna/intern/rna_nodetree.c")
    sp = SourcePath(source,"source/blender/makesrna/intern/rna_nodetree.c")
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    for x in range(0, lines):
        if "/* -- Texture Nodes --------------------------------------------------------- */" in contents[x]:
            index = x

    newline =  "static void def_cmp_{}(StructRNA *srna)\n{{\n}}\n".format(node.lower())
    contents.insert(index, newline)
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step5_registerstatictypes(source, node):
    print("Step 5: Register the node in Python RNA in /source/blender/nodes/NOD_static_types.h")
    sp = SourcePath(source,"source/blender/nodes/NOD_static_types.h")
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    for x in range(0, lines):
        if "DefNode( CompositorNode, CMP_NODE_" in contents[x]:
            index = x

    newline =  "DefNode( CompositorNode, CMP_NODE_{},".format(node.upper())
    newline = newline.ljust(50)
    newline = "{}def_cmp_{},".format(newline,node.lower())
    newline = newline.ljust(74)
    newline = "{}\"{}\",".format(newline,node.upper())
    newline = newline.ljust(92)
    newline = "{}{},".format(newline,node)
    newline = newline.ljust(110)
    newline = "{}\"{}\",".format(newline,re.sub(r'([a-z](?=[A-Z])|[A-Z](?=[A-Z][a-z]))', r'\1 ', node))
    newline = newline.ljust(131)
    newline = newline+"\"\"              )\n"
    contents.insert(index+1, newline)
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()


def step6_registernodcomposite(source, node):
    print("Step 5: Register the node in source/blender/nodes/NOD_composite.h")
    sp = SourcePath(source,"source/blender/nodes/NOD_composite.h")
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    for x in range(0, lines):
        if "void register_node_type_cmp_" in contents[x]:
            index = x

    newline =  "void register_node_type_cmp_{}(void);\n".format(node.lower())
    contents.insert(index+1, newline)
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step7_createnode(source, node):
    fn =  "source/blender/nodes/composite/nodes/node_composite_{}.c".format(node.lower())
    print("Step 7: Create a new file that defines node inputs and output in %s" % (fn))
    sp = SourcePath(source,fn)
    contents = """#include "node_composite_util.h"
    
/* **************** TKNSPACED Tool ******************** */
 
static bNodeSocketTemplate cmp_node_TKNLOWER_in[]= {
	{	SOCK_RGBA, 1, N_("Image"),			1.0f, 1.0f, 1.0f, 1.0f},
	{	SOCK_FLOAT, 1, N_("Variable"),	1.0f, 0.0f, 0.0f, 0.0f, 0.001f, 10.0f, PROP_UNSIGNED},
	{	-1, 0, ""	}
};
static bNodeSocketTemplate cmp_node_TKNLOWER_out[]= {
	{	SOCK_RGBA, 0, N_("Image")},
	{	-1, 0, ""	}
};
 
void register_node_type_cmp_TKNLOWER(void)
{
	static bNodeType ntype;
 
	node_type_base(&ntype, CMP_NODE_TKNUPPER, "TKNSPACED", NODE_CLASS_CONVERTOR, NODE_OPTIONS);
	node_type_socket_templates(&ntype, cmp_node_TKNLOWER_in, cmp_node_TKNLOWER_out);
	node_type_size(&ntype, 140, 100, 320);
 
	nodeRegisterType(&ntype);
}"""
    contents = contents.replace("TKNLOWER", node.lower())
    contents = contents.replace("TKNUPPER", node.upper())
    contents = contents.replace("TKNSPACED", re.sub(r'([a-z](?=[A-Z])|[A-Z](?=[A-Z][a-z]))', r'\1 ', node))
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()


def step8_updatecmakelists(source, node):
    print("Step 8: Make sure it is compiled by adding it to source/blender/nodes/CMakeLists.txt")
    sp = SourcePath(source,"source/blender/nodes/CMakeLists.txt")
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    for x in range(0, lines):
        if "composite/nodes/node_composite_" in contents[x]:
            index = x

    newline =  "\tcomposite/nodes/node_composite_{}.c\n".format(node.lower())
    contents.insert(index+1, newline)
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step9_createnodeheader(source, node):
    fn =  "source/blender/compositor/nodes/COM_{}Node.h".format(node)
    print("Step 9: Create the header file for the node in %s" % (fn))
    sp = SourcePath(source,fn)
    contents = """#ifndef _COM_TKNNORMALNode_h_
#define _COM_TKNNORMALNode_h_
 
#include "COM_Node.h"
 
/**
 * @brief TKNNORMALNode
 * @ingroup Node
 */
class TKNNORMALNode : public Node {
public:
	TKNNORMALNode(bNode *editorNode);
	void convertToOperations(NodeConverter &converter, const CompositorContext &context) const;
};
 
#endif"""
    contents = contents.replace("TKNNORMAL", node)
    contents = contents.replace("TKNLOWER", node.lower())
    contents = contents.replace("TKNUPPER", node.upper())
    contents = contents.replace("TKNSPACED", re.sub(r'([a-z](?=[A-Z])|[A-Z](?=[A-Z][a-z]))', r'\1 ', node))
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step10_createnode(source, node):
    fn =  "source/blender/compositor/nodes/COM_{}Node.cpp".format(node)
    print("Step 10: Create the main node file in %s" % (fn))
    sp = SourcePath(source,fn)
    contents = """#include "COM_TKNNORMALNode.h"
#include "COM_TKNNORMALOperation.h"
#include "COM_ExecutionSystem.h"
 
TKNNORMALNode::TKNNORMALNode(bNode *editorNode) : Node(editorNode)
{
	/* pass */
}
 
void TKNNORMALNode::convertToOperations(NodeConverter &converter, const CompositorContext &context) const
{
	TKNNORMALOperation *operation = new TKNNORMALOperation();

	converter.mapInputSocket(getInputSocket(0), operation->getInputSocket(0));
	converter.mapInputSocket(getInputSocket(1), operation->getInputSocket(1));
	converter.mapOutputSocket(getOutputSocket(0), operation->getOutputSocket(0));
	converter.addOperation(operation);
}"""
    contents = contents.replace("TKNNORMAL", node)
    contents = contents.replace("TKNLOWER", node.lower())
    contents = contents.replace("TKNUPPER", node.upper())
    contents = contents.replace("TKNSPACED", re.sub(r'([a-z](?=[A-Z])|[A-Z](?=[A-Z][a-z]))', r'\1 ', node))
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step11_createoperationheader(source, node):
    fn =  "source/blender/compositor/operations/COM_{}Operation.h".format(node)
    print("Step 11: Write the header file for the operation in %s" % (fn))
    sp = SourcePath(source,fn)
    contents = """#ifndef _COM_TKNNORMALOperation_h
#define _COM_TKNNORMALOperation_h
#include "COM_NodeOperation.h"
 
 
class TKNNORMALOperation : public NodeOperation {
private:
	/**
	 * Cached reference to the inputProgram
	 */
	SocketReader *m_inputProgram;
	SocketReader *m_inputTKNNORMALProgram;
 
public:
	TKNNORMALOperation();
 
	/**
	 * the inner loop of this program
	 */
	void executePixelSampled(float output[4], float x, float y, PixelSampler sampler);
 
	/**
	 * Initialize the execution
	 */
	void initExecution();
 
	/**
	 * Deinitialize the execution
	 */
	void deinitExecution();
};
#endif"""
    contents = contents.replace("TKNNORMAL", node)
    contents = contents.replace("TKNLOWER", node.lower())
    contents = contents.replace("TKNUPPER", node.upper())
    contents = contents.replace("TKNSPACED", re.sub(r'([a-z](?=[A-Z])|[A-Z](?=[A-Z][a-z]))', r'\1 ', node))
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step12_createoperation(source, node):
    fn =  "source/blender/compositor/operations/COM_{}Operation.cpp".format(node)
    print("Step 12: Create the code that does the work in %s" % (fn))
    sp = SourcePath(source,fn)
    contents = """#include "COM_TKNNORMALOperation.h"
#include "BLI_math.h"
 
TKNNORMALOperation::TKNNORMALOperation() : NodeOperation()
{
	this->addInputSocket(COM_DT_COLOR);
	this->addInputSocket(COM_DT_VALUE);
	this->addOutputSocket(COM_DT_COLOR);
	this->m_inputProgram = NULL;
	this->m_inputTKNNORMALProgram = NULL;
}
void TKNNORMALOperation::initExecution()
{
	this->m_inputProgram = this->getInputSocketReader(0);
	this->m_inputTKNNORMALProgram = this->getInputSocketReader(1);
}
 
void TKNNORMALOperation::executePixelSampled(float output[4], float x, float y, PixelSampler sampler)
{
	float inputValue[4];
	float inputTKNNORMAL[4];
 
	this->m_inputProgram->readSampled(inputValue, x, y, sampler);
	this->m_inputTKNNORMALProgram->readSampled(inputTKNNORMAL, x, y, sampler);
	const float TKNLOWER = inputTKNNORMAL[0];
	/* check for negative to avoid nan's */
	output[0] = inputValue[0] > 0.0f ? powf(inputValue[0], TKNLOWER) : inputValue[0];
	output[1] = inputValue[1] > 0.0f ? powf(inputValue[1], TKNLOWER) : inputValue[1];
	output[2] = inputValue[2] > 0.0f ? powf(inputValue[2], TKNLOWER) : inputValue[2];
 
	output[3] = inputValue[3];
}
 
void TKNNORMALOperation::deinitExecution()
{
	this->m_inputProgram = NULL;
	this->m_inputTKNNORMALProgram = NULL;
}"""
    contents = contents.replace("TKNNORMAL", node)
    contents = contents.replace("TKNLOWER", node.lower())
    contents = contents.replace("TKNUPPER", node.upper())
    contents = contents.replace("TKNSPACED", re.sub(r'([a-z](?=[A-Z])|[A-Z](?=[A-Z][a-z]))', r'\1 ', node))
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step13_updatenodeconvert(source, node):
    print("Step 13: Add the node to the converter code listing in source/blender/compositor/intern/COM_Converter.cpp")
    sp = SourcePath(source,"source/blender/compositor/intern/COM_Converter.cpp")
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    lastinclude=-1
    for x in range(0, lines):
        if "/* not inplemented yet */" in contents[x]:
            index = x
        if "#include \"COM_" in contents[x]:
            lastinclude = x        
    newline =  "\t\tcase CMP_NODE_{}:\n\t\t\tnode = new {}Node(b_node);\n\t\t\tbreak;\n".format(node.upper(),node)
    contents.insert(index+1, newline)
    newline =  "#include \"COM_{}Node.h\"\n".format(node)
    contents.insert(lastinclude+1, newline)
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step14_updatecmakelists(source, node):
    print("Step 14: Make sure it will get compiled via cmake in source/blender/compositor/CMakeLists.txt")
    sp = SourcePath(source,"source/blender/compositor/CMakeLists.txt")
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    for x in range(0, lines):
        if "\toperations/COM_" in contents[x]:
            index = x
    newline = "\n\t# TKNNORMAL Nodes\n\tnodes/COM_TKNNORMALNode.cpp\n\tnodes/COM_TKNNORMALNode.h\n\toperations/COM_TKNNORMALOperation.h\n\toperations/COM_TKNNORMALOperation.cpp\n"
    newline = newline.replace("TKNNORMAL", node)
    contents.insert(index+1, newline)
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()

def step15_nodeitems(source, node):
    print("Step 15: And finally add it to release/scripts/startup/nodeitems_builtins.py so it will show up in the UI")
    sp = SourcePath(source,"release/scripts/startup/nodeitems_builtins.py")
    f = open(sp, "r")
    contents = f.readlines()
    f.close()
    lines  = len(contents)
    index = -1
    for x in range(0, lines):
        if "CompositorNodeCategory(\"CMP_CONVERTOR\", \"Converter\", items=[" in contents[x]:
            index = x
    newline = "        NodeItem(\"CompositorNodeTKNNORMAL\"),\n"
    newline = newline.replace("TKNNORMAL", node)
    contents.insert(index+1, newline)
    f = open(sp, "w")
    contents = "".join(contents)
    f.write(contents)
    f.close()


def main(argv):
    print("Blender folder = %s" % (argv[1]))
    print("NodeName       = %s" % (argv[2]))
    print("")
    if ValidNode(argv[1], argv[2]):
        step1_updatebkenode(argv[1], argv[2])
        step2_updatenodec(argv[1], argv[2])
        step3_updatedna(argv[1], argv[2])
        step4_registerrna(argv[1], argv[2])
        step5_registerstatictypes(argv[1], argv[2])
        step6_registernodcomposite(argv[1], argv[2])
        step7_createnode(argv[1], argv[2])
        step8_updatecmakelists(argv[1], argv[2])
        step9_createnodeheader(argv[1], argv[2])
        step10_createnode(argv[1], argv[2])
        step11_createoperationheader(argv[1], argv[2])
        step12_createoperation(argv[1], argv[2])
        step13_updatenodeconvert(argv[1], argv[2])
        step14_updatecmakelists(argv[1], argv[2])
        step15_nodeitems(argv[1], argv[2])
    else:
        print("Node Already Exists, exiting...")
    pass

if __name__ == "__main__":
    main(sys.argv)