test

node_ListRangeInt.py

import bpy
from node_s import *
from util import *
from bpy.props import IntProperty, EnumProperty

'''
- range exclusive n
Start, stop, step. Like range()
Start, step, count

See class unit tests for behaviours

'''


class GenListRangeInt(Node, SverchCustomTreeNode):

    ''' Generator range list of ints '''
    bl_idname = 'GenListRangeIntNode'
    bl_label = 'List Range Int'
    bl_icon = 'OUTLINER_OB_EMPTY'

    start_ = IntProperty(
        name='start', description='start', default=0,
        options={'ANIMATABLE'}, update=updateNode)
    stop_ = IntProperty(
        name='stop', description='stop', default=10,
        options={'ANIMATABLE'}, update=updateNode)
    step_ = IntProperty(
        name='step', description='step', default=1,
        options={'ANIMATABLE'}, update=updateNode)
    count_ = IntProperty(
        name='count', description='num items', default=10,
        options={'ANIMATABLE'}, update=updateNode)

    current_mode = StringProperty(default="LAZYRANGE")

    modes = [
        ("LAZYRANGE", "Range", "Use python Range function", 1),
        ("COUNTRANGE", "Count", "Create range based on count", 2)
    ]

    def mode_change(self, context):

        # just because click doesn't mean we need to change mode
        mode = self.mode
        if mode == self.current_mode:
            return

        inputs = self.inputs
        outputs = self.outputs

        # find if input socket three has any input from another node,
        # if it does, as a nicety we keep track of it and reconnect that node
        # into the changed socket. manually disconnecting is often less work
        # than trying to find out where a connection came from.
        link = inputs[2].links[0]
        node_from = link.from_node
        socket_from = link.from_socket

        # now disconnect by removing
        outputs.remove(outputs[-1])

        if mode == 'LAZYRANGE':
            inputs.new('StringsSocket', "Range", "Range").prop_name = 'range_'
        else:
            inputs.new('StringsSocket', "Count", "Count").prop_name = 'count_'

        _output = context.nodes[node_from].outputs[socket_from]
        _input = inputs[2]
        context.node_group.links.new(_output, _input)
        self.current_mode = mode

    mode = EnumProperty(items=modes, default='LAZYRANGE', update=mode_change)

    def init(self, context):
        self.inputs.new('StringsSocket', "Start", "Start").prop_name = 'start_'
        self.inputs.new('StringsSocket', "Stop", "Stop").prop_name = 'stop_'
        self.inputs.new('StringsSocket', "Step", "Step").prop_name = 'step_'

        self.outputs.new('StringsSocket', "Range", "Range").prop_name = 'range_'

    def draw_buttons(self, context, layout):
        layout.prop(self, "mode", expand=True)
        # layout.prop(self, "start_", text="start")
        # layout.prop(self, "stop_", text="stop")

        # if self.mode == "LAZYRANGE":
        #     layout.prop(self, "step_", text="step")
        # else:
        #     layout.prop(self, "count_", text="count")

    def update(self):
        inputs = self.inputs
        outputs = self.outputs

        # outputs, end early.
        if not 'Range' in self.outputs or not (len(self.outputs['Rage'].links) > 0):
            return

        # inputs, both modes use these
        if 'Start' in self.inputs and self.inputs['Start'].links:
            tmp = SvGetSocketAnyType(self, self.inputs['Start'])
            Start = tmp[0][0]
        else:
            Start = self.start_

        if 'Step' in self.inputs and self.inputs['Step'].links:
            tmp = SvGetSocketAnyType(self, self.inputs['Step'])
            Step = tmp[0][0]
        else:
            Step = self.step_

        # mode dependant stuff
        if self.mode == 'LAZYRANGE':
            if 'Stop' in self.inputs and self.inputs['Stop'].links:
                tmp = SvGetSocketAnyType(self, self.inputs['Stop'])
                Stop = tmp[0][0]
            else:
                Stop = self.stop_

            range_ = [c for c in self.intRange(Start, Stop, Step)]

        # mode dependant stuff
        elif self.mode == 'COUNTRANGE':
            if 'Count' in self.inputs and self.inputs['Count'].links:
                tmp = SvGetSocketAnyType(self, self.inputs['Count'])
                Count = tmp[0][0]
            else:
                Count = self.count_

        SvSetSocketAnyType(self, 'Range', [range_])

        def intRange(self, start=0, stop=1, step=1):
            '''
            slightly different behaviour: "lazy range"
            - step is always |step| (absolute)
            - step is converted to negative if stop is less than start
            '''
            if start == stop:
                return []

            step = max(step, 1)
            if stop < start:
                step *= -1
            return list(range(start, stop, step))

        def countRange(self, start, step, count):
            count = max(count, 0)
            if count == 0:
                return []

            stop = (count*step) + start
            return list(range(start, stop, step))

        def print_tests(self, tests):
            for i in tests:
                print(i, ':', eval(i))

        def unit_tests(self):
            # test intRange "lazy range"
            a = 'intRange(20, 40, 2)'
            b = 'intRange(20, 30, 1)'
            c = 'intRange(-4, 4, 1)'
            d = 'intRange(5, -4, 1)'
            e = 'intRange(20, 30, -1)'
            self.print_tests([a, b, c, d, e])

            # test intRange "lazy range"
            a = 'countRange(20, 1, 3)'
            b = 'countRange(20, 2, 8)'
            c = 'countRange(-4, -1, 7)'
            d = 'countRange(5, -4, 1)'
            e = 'countRange(20, 30, -1)'
            self.print_tests([a, b, c, d, e])

    def update_socket(self, context):
        self.update()


def register():
    bpy.utils.register_class(GenListRangeInt)


def unregister():
    bpy.utils.unregister_class(GenListRangeInt)