PDG/partitionbyparity.py

#
# Copyright (c) 2024
#       Side Effects Software Inc.  All rights reserved.
#
# Redistribution and use of Houdini Development Kit samples in source and
# binary forms, with or without modification, are permitted provided that the
# following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. The name of Side Effects Software may not be used to endorse or
#    promote products derived from this software without specific prior
#    written permission.
#
# THIS SOFTWARE IS PROVIDED BY SIDE EFFECTS SOFTWARE `AS IS' AND ANY EXPRESS
# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN
# NO EVENT SHALL SIDE EFFECTS SOFTWARE BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#----------------------------------------------------------------------------
# This script provides an example of writing a custom PDG node in Python
#

"""This is the Python equivalent of PDG_PartitionByParity.C which uses
the C++ API.

To use this code:
    1) Create a directory on the Houdini search path for custom PDG Python
    nodes, if one does not already exist. For example, ~/houdini20.0/pdg/types
    2) Save this file to that directory as partitionbyparity.py
    3) Within Houdini the custom node will now show up in the TAB menu in
    a TOP network
"""

import pdg

from pdg.partitioner import PyPartitioner

class PartitionByParity(PyPartitioner):
    """
    Partitioner implementation that groups upstream work items based on the
    parity of an attribute value
    """
    def __init__(self, node):
        """
        __init__(self, pdg.Node) -> NoneType

        Initializes the partitioner
        """
        PyPartitioner.__init__(self, node)

    @classmethod
    def templateName(cls):
        """
        templateName(cls) -> str
        """
        return "partitionbyparity"

    @classmethod
    def templateBody(cls):
        return """{
                    "name": "partitionbyparity",
                    "parameters" : [
                        {
                            "name" : "attributename",
                            "label" : "Attribute Name",
                            "type" : "String",
                            "value" : ""
                        }
                    ]
                }"""

    def onPartition(self, partitions, work_items):
        attribute_name = self['attributename'].evaluateString()

        for work_item in work_items:
            attrib = work_item.attrib(attribute_name)
            if not attrib:
                self.cookWarning(
                    "'{}' has no attribute with the name '{}'",
                    work_item.name, attribute_name)
                continue

            index = 0 if int(attrib.asNumber(0)) % 2 == 0 else 1
            partitions.addItemToPartition(work_item, index)

        return pdg.result.Success

def registerTypes(type_registry):
    type_registry.registerNode(
            PartitionByParity,
            pdg.nodeType.Partitioner,
            label="Partition by Parity")