Maya USD Import Chaser in Python to import `primvars:cbId_transform` attribute to the transform as `cbId` instead of as attribute on the shape

maya_usd_import_chaser_cbId_to_transform.py

import logging

import mayaUsd.lib as mayaUsdLib
from maya import cmds
import maya.api.OpenMaya as OpenMaya


def create_cbid_attr(default_value):
    default = OpenMaya.MFnStringData().create(default_value)
    return OpenMaya.MFnTypedAttribute().create("cbId",
                                               "cbId",
                                               OpenMaya.MFnData.kString,
                                               default)


class ImportChaserCbIdShapes(mayaUsdLib.ImportChaser):
    """Import `cbId_transform` attribute to `cbId` on transform node"""
    def __init__(self, factoryContext, *args, **kwargs):
        super(ImportChaserCbIdShapes, self).__init__(factoryContext,
                                                     *args,
                                                     **kwargs)
        self.modifier = OpenMaya.MDGModifier()

    def _post_import(self, returnPredicate, stage, dagPath, sdfPath, jobArgs):
        
        fn = OpenMaya.MFnDependencyNode()
        # Note: The Sdf to Dag map only contains the dag entries for which
        #  an SdfPath exists. Since shapes are expanded from its parent
        #  they do not get their own mapping. As such, we need to manually
        #  detect whether we're dealing with a shape currently
        for mapping in self.GetSdfToDagMap():
            dag_path = mapping.data()
            sdf_path = mapping.key()
            
            prim = stage.GetPrimAtPath(sdf_path)
            if not prim.HasAttribute("primvars:cbId_transform"):
                continue
                
            num = dag_path.numberOfShapesDirectlyBelow()
            if num == 0:
                continue
                
            if num > 1:
                logger.warning("More than one shape for: {}".format(dag_path))
            
            # Get the shape
            shape_dag_path = OpenMaya.MDagPath(dag_path).extendToShape(0)
            shape_node = shape_dag_path.node()
            fn.setObject(shape_node)
            
            # Remove cbId_transform from shape if it exists tehre
            plug = fn.hasAttribute("cbId_transform")
            if plug:
                shape_attr = fn.attribute("cbId_transform")
                self.modifier.removeAttribute(shape_node, shape_attr)
                
            # Add to transform instead
            value = prim.GetAttribute("primvars:cbId_transform").Get() or ""
            attr = create_cbid_attr(value)
            
            transform_node = dag_path.node()
            self.modifier.addAttribute(transform_node, attr)

        self.modifier.doIt()

        return True

    def PostImport(self, returnPredicate, stage, dagPaths, sdfPaths, jobArgs):
        try:
            return self._post_import(
                returnPredicate,
                stage,
                dagPaths,
                sdfPaths,
                jobArgs
            )
        except Exception as exc:
            logging.error(exc, exc_info=True)
            raise

    def Redo(self):
        # NOTE: (yliangsiew) Undo the undo to re-do.
        self.modifier.doIt()
        return True

    def Undo(self):
        self.modifier.undoIt()
        return True


mayaUsdLib.ImportChaser.Register(ImportChaserCbIdShapes, "cbId_transform")
path = r"P:\Projects\test\production\usd_cbId_test\work\FX\houdini\geo\with_suffix_roy.usda"
nodes = cmds.file(path,
                  sharedReferenceFile=False,
                  reference=True,
                  returnNewNodes=True,
                  options="chaser=[cbId_transform]")

#rootPaths = cmds.mayaUSDImport(v=True,
#                               f=path,
#                               chaser=['cbId_transform'])

Maya USD export chaser to write primvars:cbId_transform and primvars:cbId attributes

import logging

from maya.api import OpenMaya
import mayaUsd.lib as mayaUsdLib
from pxr import Sdf


class CbExportChaser(mayaUsdLib.ExportChaser):

    def __init__(self, factoryContext, *args, **kwargs):
        super(CbExportChaser, self).__init__(factoryContext, *args, **kwargs)
        self.dag_to_usd = factoryContext.GetDagToUsdMap()
        self.stage = factoryContext.GetStage()
        self.job_args = factoryContext.GetJobArgs()

    def PostExport(self):
        try:
            return self._post_export()
        except Exception as exc:
            logging.error(exc, exc_info=True)
            raise

    def _post_export(self):
        
        merged_shapes = getattr(self.job_args, "mergeTransformAndShape", True)
        
        fn = OpenMaya.MFnDependencyNode()
        for entry in self.dag_to_usd:
            prim_path = entry.data()
            dag_path = entry.key()

            fn.setObject(dag_path.node())
            if not fn.hasAttribute("cbId"):
                continue

            try:
                value = fn.findPlug("cbId", True).asString()
            except RuntimeError:
                logging.error(
                    "Unable to get {dag_path}.cbId value.".format(
                        dag_path=dag_path
                    ), 
                    exc_info=True
                )
                continue

            # Create attributes
            is_parent_of_shape = merged_shapes and bool(dag_path.numberOfShapesDirectlyBelow())
            attr_name = "cbId_transform" if is_parent_of_shape else "cbId" 
            attr_name = "primvars:{}".format(attr_name)

            prim = self.stage.GetPrimAtPath(prim_path)
            attr = prim.CreateAttribute(attr_name, Sdf.ValueTypeNames.String)
            attr.Set(value)

        return True


# Testing
from maya import cmds

mayaUsdLib.ExportChaser.Unregister(CbExportChaser, "cbid")
mayaUsdLib.ExportChaser.Register(CbExportChaser, "cbid")

cmds.loadPlugin("mayaUsdPlugin", quiet=True)
cmds.mayaUSDExport(file="E:/test.usda",
                   frameRange=(1, 5),
                   frameStride=1.0,
                   chaser=["cbid"],
                   # mergeTransformAndShape=True,
                   selection=True)