<commit_before>import collections
import pymel.core as pymel
from omtk.core.classCtrl import BaseCtrl
from omtk.core.classModule import Module
from omtk.libs import libRigging, libCtrlShapes
from omtk.core import consts_omtk

class CtrlFk(BaseCtrl):
    def __createNode__(self, *args, **kwargs):
        '''
        if 'shoulder' in name.lower():
            node = libCtrlShapes.create_shape_double_needle(size=size*0.04, normal=(0, 0, 1), *args, **kwargs)
        else:
        '''
        node = super(CtrlFk, self).__createNode__(multiplier=1.1, *args, **kwargs)

        make = next(iter(node.inputs()), None)
        if make:
            # TODO: Multiply radius???
            #make.radius.set(size)
            make.degree.set(1)
            make.sections.set(8)

        return node


class FK(Module):
    DEFAULT_NAME_USE_FIRST_INPUT = True

    def __init__(self, *args, **kwargs):
        super(FK, self).__init__(*args, **kwargs)
        self.ctrls = None
        self.sw_translate=False
        self.create_spaceswitch = True

    #
    # libSerialization implementation
    #
    def __callbackNetworkPostBuild__(self):
        """
        Cleaning routine automatically called by libSerialization after a network import.
        """
        # Ensure there's no None value in the .ctrls array.
        # This can happen if the rigging delete the stored shape before rebuilding.
        try:
            self.ctrls= filter(None, self.ctrls)
        except (AttributeError, TypeError):
            pass
        super(FK, self).__callbackNetworkPostBuild__()

    def build(self, rig, constraint=True, parent=True, *args, **kwargs):
        super(FK, self).build(rig, create_grp_rig=False, *args, **kwargs)

        nomenclature_anm = self.get_nomenclature_anm(rig)

        # Define ctrls
        if not self.ctrls:
            self.ctrls = []
            for input in self.chain_jnt:
                ctrl = CtrlFk()
                self.ctrls.append(ctrl)

        # Create ctrls
        for input, ctrl in zip(self.chain_jnt, self.ctrls):
            ctrl_nomenclature = nomenclature_anm.rebuild(input.name())
            ctrl_name = ctrl_nomenclature.resolve('fk')
            ctrl.build(rig, name=ctrl_name, refs=input)
            ctrl.setMatrix(input.getMatrix(worldSpace=True))

        if self.create_spaceswitch:
            if self.sw_translate:
                self.ctrls[0].create_spaceswitch(rig, self, self.parent, add_world=True)
            else:
                self.ctrls[0].create_spaceswitch(rig, self, self.parent, skipTranslate=['x', 'y', 'z'], add_world=True)

        self.ctrls[0].setParent(self.grp_anm)
        for i in range(1, len(self.ctrls)):
            self.ctrls[i].setParent(self.ctrls[i - 1])

        # Connect jnt -> anm
        if constraint is True:
            for inn, ctrl in zip(self.chain_jnt, self.ctrls):
                pymel.parentConstraint(ctrl, inn)
                pymel.connectAttr(ctrl.scaleX, inn.scaleX)
                pymel.connectAttr(ctrl.scaleY, inn.scaleY)
                pymel.connectAttr(ctrl.scaleZ, inn.scaleZ)

        '''
        # Connect to parent
        if parent and self.parent is not None:
            pymel.parentConstraint(self.parent, self.grp_anm, maintainOffset=True)
            #pymel.scaleConstraint(self.parent, self.grp_anm, maintainOffset=True)
        '''

    def unbuild(self):
        super(FK, self).unbuild()


class CtrlFkAdd(BaseCtrl):
    def __createNode__(self, size=None, refs=None, *args, **kwargs):
        # Resolve size automatically if refs are provided.
        ref = next(iter(refs), None) if isinstance(refs, collections.Iterable) else refs
        if size is None and ref is not None:
            size = libRigging.get_recommended_ctrl_size(ref)
        else:
            size = 1.0

        node = libCtrlShapes.create_shape_needle(size=size, *args, **kwargs)
<commit_msg>Fk module now support multiple fk chain<commit_after>import collections
import pymel.core as pymel
from omtk.core.classCtrl import BaseCtrl
from omtk.core.classModule import Module
from omtk.libs import libRigging, libCtrlShapes
from omtk.core import consts_omtk

class CtrlFk(BaseCtrl):
    def __createNode__(self, *args, **kwargs):
        '''
        if 'shoulder' in name.lower():
            node = libCtrlShapes.create_shape_double_needle(size=size*0.04, normal=(0, 0, 1), *args, **kwargs)
        else:
        '''
        node = super(CtrlFk, self).__createNode__(multiplier=1.1, *args, **kwargs)

        make = next(iter(node.inputs()), None)
        if make:
            # TODO: Multiply radius???
            #make.radius.set(size)
            make.degree.set(1)
            make.sections.set(8)

        return node


class FK(Module):
    DEFAULT_NAME_USE_FIRST_INPUT = True

    def __init__(self, *args, **kwargs):
        super(FK, self).__init__(*args, **kwargs)
        self.ctrls = []
        self.sw_translate=False
        self.create_spaceswitch = True

    #
    # libSerialization implementation
    #
    def __callbackNetworkPostBuild__(self):
        """
        Cleaning routine automatically called by libSerialization after a network import.
        """
        # Ensure there's no None value in the .ctrls array.
        # This can happen if the rigging delete the stored shape before rebuilding.
        try:
            self.ctrls= filter(None, self.ctrls)
        except (AttributeError, TypeError):
            pass
        super(FK, self).__callbackNetworkPostBuild__()

    def build(self, rig, constraint=True, parent=True, create_grp_anm=True, *args, **kwargs):
        super(FK, self).build(rig, create_grp_rig=False, *args, **kwargs)

        nomenclature_anm = self.get_nomenclature_anm(rig)

        # Define ctrls
        num_ctrls = 0
        if self.ctrls:
            num_ctrls = len(self.ctrls)
        chain_first_ctrl_idx = 0
        for i, chain in enumerate(self.chains):
            if not self.ctrls or chain_first_ctrl_idx + len(chain) > num_ctrls:
                for input in chain:
                    ctrl = CtrlFk()
                    self.ctrls.append(ctrl)
                    num_ctrls += 1

            # Create ctrls
            for input, ctrl in zip(chain, self.ctrls[chain_first_ctrl_idx:chain_first_ctrl_idx + len(chain)]):
                ctrl_nomenclature = nomenclature_anm.rebuild(input.name())
                ctrl_name = ctrl_nomenclature.resolve('fk')
                ctrl.build(rig, name=ctrl_name, refs=input)
                ctrl.setMatrix(input.getMatrix(worldSpace=True))

            if self.create_spaceswitch:
                if self.sw_translate:
                    self.ctrls[chain_first_ctrl_idx].create_spaceswitch(rig, self, self.parent, add_world=True)
                else:
                    self.ctrls[chain_first_ctrl_idx].create_spaceswitch(rig, self, self.parent,
                                                                        skipTranslate=['x', 'y', 'z'], add_world=True)

            self.ctrls[chain_first_ctrl_idx].setParent(self.grp_anm)
            for j in range(chain_first_ctrl_idx + 1, chain_first_ctrl_idx + len(chain)):
                self.ctrls[j].setParent(self.ctrls[j - 1])

            # Connect jnt -> anm
            if constraint is True:
                for inn, ctrl in zip(chain, self.ctrls[chain_first_ctrl_idx:chain_first_ctrl_idx + len(chain)]):
                    pymel.parentConstraint(ctrl, inn)
                    pymel.connectAttr(ctrl.scaleX, inn.scaleX)
                    pymel.connectAttr(ctrl.scaleY, inn.scaleY)
                    pymel.connectAttr(ctrl.scaleZ, inn.scaleZ)

            chain_first_ctrl_idx += len(chain)

        '''
        # Connect to parent
        if parent and self.parent is not None:
            pymel.parentConstraint(self.parent, self.grp_anm, maintainOffset=True)
            #pymel.scaleConstraint(self.parent, self.grp_anm, maintainOffset=True)
        '''

    def unbuild(self):
        super(FK, self).unbuild()


class CtrlFkAdd(BaseCtrl):
    def __createNode__(self, size=None, refs=None, *args, **kwargs):
        # Resolve size automatically if refs are provided.
        ref = next(iter(refs), None) if isinstance(refs, collections.Iterable) else refs
        if size is None and ref is not None:
            size = libRigging.get_recommended_ctrl_size(ref)
        else:
            size = 1.0

        node = libCtrlShapes.create_shape_needle(size=size, *args, **kwargs)

{'fraction_non_alphanumeric': 0.10222169699834555, 'fraction_numerical': 0.004017962656582368, 'mean_word_length': 2.2650462962962963, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '791302', 'n_tokens_mistral': 2656, 'n_tokens_neox': 2569, 'n_words': 568}