EMSUSD-254 Allow selecting the referenced prim

lib/mayaUsd/resources/scripts/mayaUsdLibRegisterStrings.py

@@ -73,6 +73,15 @@ def mayaUsdLibRegisterStrings():
     register('kTextVariant', 'Variant')
     register('kTextVariantToolTip','If selected, your Maya reference will be defined in a variant. This will enable your prim to\nhave 2 variants you can switch between in the Outliner; the Maya reference and its USD cache.')
 
+    register('kAddRefOrPayloadPrimPathToolTip', 'Specifying a prim path will make an explicit reference to a prim.\n' +
+                                        'If there is no default prim in your chosen reference, this field auto populates with the top level prim for you to reference in.\n' +
+                                        'If the field is left blank, no prim will be referenced.')
+    register('kAddRefOrPayloadPrimPathLabel', 'Prim Path')
+    register('kAddRefOrPayloadPrimPathPlaceHolder', ' (Default Prim)')
+    register('kAddRefOrPayloadPrimPathHelpLabel', 'Help on Select a Prim for Reference')
+    register('kAddRefOrPayloadPrimPathTitle', 'Select a Prim to Reference')
+    register('kAddRefOrPayloadSelectLabel', 'Select')
+
     # mayaUsdClearRefsOrPayloadsOptions.py
     register('kClearRefsOrPayloadsOptionsTitle', 'Clear All USD References/Payloads')
     register('kClearRefsOrPayloadsOptionsMessage', 'Clear all references/payloads on %s?')

lib/mayaUsd/resources/scripts/mayaUsdMayaReferenceUtils.py

@@ -16,6 +16,7 @@
 
 import maya.cmds as cmds
 from mayaUsdLibRegisterStrings import getMayaUsdLibString
+from pxr import Sdf
 
 # These names should not be localized as Usd only accepts [a-z,A-Z] as valid characters.
 kDefaultMayaReferencePrimName = 'MayaReference1'
@@ -27,13 +28,15 @@
 compositionArcReference = 'Reference'
 _compositionArcValues = [compositionArcReference, compositionArcPayload]
 
-listEditTypeKey = ''
+listEditTypeKey = 'listEditType'
 listEditTypePrepend = 'Prepend'
 listEditTypeAppend = 'Append'
 _listEditedAsValues = [listEditTypePrepend, listEditTypeAppend]
 
 loadPayloadKey = 'loadPayload'
 
+referencedPrimPathKey = 'referencedPrimPath'
+
 def defaultMayaReferencePrimName():
     return kDefaultMayaReferencePrimName
 
@@ -118,7 +121,84 @@ def _compositionArcChanged(selectedItem):
     enableLoadPayload = bool(compositionArc == compositionArcPayload)
     cmds.checkBoxGrp('loadPayload', edit=True, enable=enableLoadPayload)
 
-def createUsdRefOrPayloadUI(showLoadPayload=False):
+def _selectReferencedPrim(*args):
+    """
+    Open a dialog to select a prim from the currently selected file
+    to be the referenced prim.
+    """
+    filename = _getCurrentFilename()
+    primPath = cmds.textFieldGrp('mayaUsdAddRefOrPayloadPrimPath', query=True, text=True)
+    title = getMayaUsdLibString('kAddRefOrPayloadPrimPathTitle')
+    helpLabel = getMayaUsdLibString('kAddRefOrPayloadPrimPathHelpLabel')
+    helpToken = 'something something'
+    result = cmds.usdImportDialog(filename, hideVariants=True, hideRoot=True, primPath=primPath, title=title, helpLabel=helpLabel, helpToken=helpToken)
+    if result:
+        primPath = cmds.usdImportDialog(query=True, primPath=True)
+        cmds.textFieldGrp('mayaUsdAddRefOrPayloadPrimPath', edit=True, text=primPath)
+    cmds.usdImportDialog(clearData=True)
+
+_currentFileName = None
+
+def _setCurrentFilename(filename):
+    """Sets the current file selection."""
+    global _currentFileName
+    _currentFileName = filename
+
+def _getCurrentFilename():
+    """Retrieve the current file selection."""
+    global _currentFileName
+    return _currentFileName
+
+def _resetReferencedPrim(*args):
+    """Reset the referenced prim UI"""
+    _updateReferencedPrimBasedOnFile()
+
+def _getDefaultAndRootPrims(filename):
+    """Retrieve the default and first root prims of a USD file."""
+    defPrim, rootPrim = None, None
+    try:
+        layer = Sdf.Layer.FindOrOpen(filename)
+        if layer:
+            # Note: the root prims at the USD layer level are SdfPrimSpec,
+            #       so they are not SdfPath themselves nor prim. That is
+            #       why their path is retrieved via their path property.
+            #
+            #       The default prim is a pure token though, because it is
+            #       a metadata on the layer, so it can be used as-is.
+            rootPrims = layer.rootPrims
+            rootPrim = rootPrims[0].path if len(rootPrims) > 0 else None
+            defPrim = layer.defaultPrim
+    except Exception as ex:
+        print(str(ex))
+
+    return defPrim, rootPrim
+
+def _updateReferencedPrimBasedOnFile():
+    """Update all UI related to the referenced prim based on the currently selected file."""
+    filename = _getCurrentFilename()
+    if not filename:
+        cmds.textFieldGrp('mayaUsdAddRefOrPayloadPrimPath', edit=True, text='', placeholderText='')
+        cmds.textFieldGrp('mayaUsdAddRefOrPayloadPrimPath', edit=True, enable=False)
+        cmds.button('mayaUsdAddRefOrPayloadFilePathBrowser', edit=True, enable=False)
+        cmds.symbolButton('mayaUsdAddRefOrPayloadFilePathReset', edit=True, enable=False)
+        return
+
+    cmds.textFieldGrp('mayaUsdAddRefOrPayloadPrimPath', edit=True, enable=True)
+    cmds.button('mayaUsdAddRefOrPayloadFilePathBrowser', edit=True, enable=True)
+    cmds.symbolButton('mayaUsdAddRefOrPayloadFilePathReset', edit=True, enable=True)
+
+    defaultPrim, rootPrim = _getDefaultAndRootPrims(filename)
+    if defaultPrim:
+        placeHolder = defaultPrim + getMayaUsdLibString('kAddRefOrPayloadPrimPathPlaceHolder')
+        cmds.textFieldGrp('mayaUsdAddRefOrPayloadPrimPath', edit=True, text='', placeholderText=placeHolder)
+    elif rootPrim:
+        cmds.textFieldGrp('mayaUsdAddRefOrPayloadPrimPath', edit=True, text=rootPrim, placeholderText='')
+    else:
+        cmds.textFieldGrp('mayaUsdAddRefOrPayloadPrimPath', edit=True, text='', placeholderText='')
+
+def createUsdRefOrPayloadUI(uiForLoad=False):
+    _setCurrentFilename(None)
+
     with SetParentContext(cmds.rowLayout(numberOfColumns=2)):
         tooltip = getMayaUsdLibString('kOptionAsUSDReferenceToolTip')
         cmds.optionMenuGrp('compositionArcTypeMenu',
@@ -136,27 +216,48 @@ def createUsdRefOrPayloadUI(showLoadPayload=False):
         for label in listEditedAsLabels:
             cmds.menuItem(label=label)
 
-    if showLoadPayload:
+    if uiForLoad:
+        with SetParentContext(cmds.rowLayout(numberOfColumns=3, adjustableColumn1=True)):
+            tooltip = getMayaUsdLibString('kAddRefOrPayloadPrimPathToolTip')
+            primPathLabel = getMayaUsdLibString("kAddRefOrPayloadPrimPathLabel")
+            selectLabel = getMayaUsdLibString("kAddRefOrPayloadSelectLabel")
+            cmds.textFieldGrp('mayaUsdAddRefOrPayloadPrimPath', label=primPathLabel, ann=tooltip)
+            cmds.button('mayaUsdAddRefOrPayloadFilePathBrowser', label=selectLabel, ann=tooltip)
+            cmds.symbolButton('mayaUsdAddRefOrPayloadFilePathReset', image="refresh.png", ann=tooltip)
+
         cmds.checkBoxGrp('loadPayload',
                          label=getMayaUsdLibString('kOptionLoadPayload'),
                          annotation=getMayaUsdLibString('kLoadPayloadAnnotation'),
                          ncb=1)
 
-def initUsdRefOrPayloadUI(values, showLoadPayload=False):
+def initUsdRefOrPayloadUI(values, uiForLoad=False):
     compositionArcMenuIndex = _getMenuIndex(_compositionArcValues, values[compositionArcKey])
     cmds.optionMenuGrp('compositionArcTypeMenu', edit=True, select=compositionArcMenuIndex)
 
     listEditTypeMenuIndex = _getMenuIndex(_listEditedAsValues,values[listEditTypeKey])
     cmds.optionMenu('listEditedAsMenu', edit=True, select=listEditTypeMenuIndex)
 
-    if showLoadPayload:
+    if uiForLoad:
         cmds.optionMenuGrp('compositionArcTypeMenu', edit=True, cc=_compositionArcChanged)
         _compositionArcChanged(compositionArcMenuIndex)
         cmds.checkBoxGrp('loadPayload', edit=True, value1=values[loadPayloadKey])
 
-def commitUsdRefOrPayloadUI(showLoadPayload=False):
+        cmds.button("mayaUsdAddRefOrPayloadFilePathBrowser", c=_selectReferencedPrim, edit=True)
+        cmds.button("mayaUsdAddRefOrPayloadFilePathReset", c=_resetReferencedPrim, edit=True)
+
+        _updateReferencedPrimBasedOnFile()
+
+def updateUsdRefOrPayloadUI(selectedFile):
+    if selectedFile == _getCurrentFilename():
+        return
+    _setCurrentFilename(selectedFile)
+    _updateReferencedPrimBasedOnFile()
+
+def commitUsdRefOrPayloadUI(uiForLoad=False):
     values = {}
     values[compositionArcKey] = _getMenuGrpValue('compositionArcTypeMenu', _compositionArcValues)
     values[listEditTypeKey  ] = _getMenuValue('listEditedAsMenu', _listEditedAsValues)
-    values[loadPayloadKey   ] = cmds.checkBoxGrp('loadPayload', query=True, value1=True) if showLoadPayload else True
+    values[loadPayloadKey   ] = cmds.checkBoxGrp('loadPayload', query=True, value1=True) if uiForLoad else True
+    values[referencedPrimPathKey] = cmds.textFieldGrp('mayaUsdAddRefOrPayloadPrimPath', query=True, text=True) if uiForLoad else ''
+
     return values

lib/mayaUsd/resources/scripts/mayaUsdUtils.py

@@ -211,6 +211,16 @@ def saveWantPayloadLoaded(want):
     opVarName = "mayaUsd_WantPayloadLoaded"
     cmds.optionVar(iv=(opVarName, want))
 
+def getReferencedPrimPath():
+    opVarName = "mayaUsd_ReferencedPrimPath"
+    if not cmds.optionVar(exists=opVarName):
+        return ''
+    return cmds.optionVar(query=opVarName)
+
+def saveReferencedPrimPath(primPath):
+    opVarName = "mayaUsd_ReferencedPrimPath"
+    cmds.optionVar(sv=(opVarName, primPath))
+
 def showHelpMayaUSD(contentId):
     """
     Helper method to display help content.

lib/mayaUsd/ufe/MayaUsdContextOps.cpp

@@ -184,6 +184,7 @@ const char* _selectUSDFileScript()
             string $result[] = `fileDialog2
                 -fileMode 1
                 -caption "Add USD Reference/Payload to Prim"
+                -okCaption Reference
                 -fileFilter "USD Files (%s);;%s"
                 -optionsUICreate addUSDReferenceCreateUi
                 -optionsUIInit addUSDReferenceInitUi
@@ -555,10 +556,12 @@ Ufe::UndoableCommand::Ptr MayaUsdContextOps::doOpCmd(const ItemPath& itemPath)
         if (path.empty())
             return nullptr;
 
-        const bool asRef = UsdMayaUtilFileSystem::wantReferenceCompositionArc();
-        const bool prepend = UsdMayaUtilFileSystem::wantPrependCompositionArc();
+        const std::string primPath = UsdMayaUtilFileSystem::getReferencedPrimPath();
+        const bool        asRef = UsdMayaUtilFileSystem::wantReferenceCompositionArc();
+        const bool        prepend = UsdMayaUtilFileSystem::wantPrependCompositionArc();
         if (asRef) {
-            return std::make_shared<UsdUfe::UsdUndoAddReferenceCommand>(prim(), path, prepend);
+            return std::make_shared<UsdUfe::UsdUndoAddReferenceCommand>(
+                prim(), path, primPath, prepend);
         } else {
             Ufe::UndoableCommand::Ptr preloadCmd;
             const bool                preload = UsdMayaUtilFileSystem::wantPayloadLoaded();
@@ -569,8 +572,8 @@ Ufe::UndoableCommand::Ptr MayaUsdContextOps::doOpCmd(const ItemPath& itemPath)
                 preloadCmd = std::make_shared<UsdUfe::UsdUndoUnloadPayloadCommand>(prim());
             }
 
-            auto payloadCmd
-                = std::make_shared<UsdUfe::UsdUndoAddPayloadCommand>(prim(), path, prepend);
+            auto payloadCmd = std::make_shared<UsdUfe::UsdUndoAddPayloadCommand>(
+                prim(), path, primPath, prepend);
 
             auto compoCmd = std::make_shared<Ufe::CompositeUndoableCommand>();
             compoCmd->append(preloadCmd);

lib/mayaUsd/utils/utilFileSystem.cpp

@@ -422,6 +422,14 @@ bool UsdMayaUtilFileSystem::wantPayloadLoaded()
         && MGlobal::optionVarIntValue(WANT_PAYLOAD_LOADED);
 }
 
+std::string UsdMayaUtilFileSystem::getReferencedPrimPath()
+{
+    static const MString WANT_REFERENCE_COMPOSITION_ARC = "mayaUsd_ReferencedPrimPath";
+    if (!MGlobal::optionVarExists(WANT_REFERENCE_COMPOSITION_ARC))
+        return {};
+    return MGlobal::optionVarStringValue(WANT_REFERENCE_COMPOSITION_ARC).asChar();
+}
+
 const char* getScenesFolderScript = R"(
 global proc string UsdMayaUtilFileSystem_GetScenesFolder()
 {

lib/mayaUsd/utils/utilFileSystem.h

@@ -171,6 +171,11 @@ bool wantPrependCompositionArc();
 MAYAUSD_CORE_PUBLIC
 bool wantPayloadLoaded();
 
+/*! \brief returns the prim path referenced by the USD reference or payload.
+ */
+MAYAUSD_CORE_PUBLIC
+std::string getReferencedPrimPath();
+
 /*! \brief prepares the UI used to save layers, so that the UI can potentially make the
            selected file name relative to the given directory.
  */

lib/usd/ui/importDialog/ItemDelegate.cpp

@@ -85,7 +85,11 @@ void ItemDelegate::paint(
     // 	painter->drawRect(QRect(0, option.rect.y(), option.rect.right(), option.rect.bottom()));
     // 	painter->restore();
 
-    ParentClass::paint(painter, option, index);
+    QStyleOptionViewItem adjustedOption(option);
+    adjustedOption.decorationPosition = QStyleOptionViewItem::Right;
+    adjustedOption.decorationAlignment = Qt::AlignLeft | Qt::AlignVCenter;
+
+    ParentClass::paint(painter, adjustedOption, index);
 }
 
 void ItemDelegate::setEditorData(QWidget* editor, const QModelIndex& index) const

lib/usd/ui/importDialog/TreeItem.cpp

@@ -25,44 +25,53 @@
 
 namespace MAYAUSD_NS_DEF {
 
-QPixmap* TreeItem::fsCheckBoxOn = nullptr;
-QPixmap* TreeItem::fsCheckBoxOnDisabled = nullptr;
-QPixmap* TreeItem::fsCheckBoxOff = nullptr;
-QPixmap* TreeItem::fsCheckBoxOffDisabled = nullptr;
+const QPixmap* TreeItem::fsCheckBoxOn = nullptr;
+const QPixmap* TreeItem::fsCheckBoxOnDisabled = nullptr;
+const QPixmap* TreeItem::fsCheckBoxOff = nullptr;
+const QPixmap* TreeItem::fsCheckBoxOffDisabled = nullptr;
 
-TreeItem::TreeItem(const UsdPrim& prim, Type t) noexcept
+TreeItem::TreeItem(const UsdPrim& prim, bool isDefaultPrim, Column column) noexcept
     : ParentClass()
     , fPrim(prim)
-    , fType(t)
+    , fColumn(column)
     , fCheckState(CheckState::kChecked_Disabled)
     , fVariantSelectionModified(false)
 {
-    initializeItem();
+    initializeItem(isDefaultPrim);
 }
 
 UsdPrim TreeItem::prim() const { return fPrim; }
 
 int TreeItem::type() const { return ParentClass::ItemType::UserType; }
 
+const QPixmap* TreeItem::createPixmap(const char* pixmapURL) const
+{
+    const QPixmap* pixmap = nullptr;
+
+    const TreeModel* treeModel = qobject_cast<const TreeModel*>(model());
+    if (treeModel != nullptr) {
+        pixmap = treeModel->mayaQtUtil().createPixmap(pixmapURL);
+    } else {
+        // The tree model should never be null, but we can recover here if it is.
+        TF_RUNTIME_ERROR("Unexpected null tree model");
+        pixmap = new QPixmap(pixmapURL);
+    }
+
+    // If the resource fails to load, return a non-null pointer.
+    static const QPixmap empty;
+    if (!pixmap)
+        pixmap = &empty;
+
+    return pixmap;
+}
+
 const QPixmap& TreeItem::checkImage() const
 {
     if (fsCheckBoxOn == nullptr) {
-        const TreeModel* treeModel = qobject_cast<const TreeModel*>(model());
-        if (treeModel != nullptr) {
-            fsCheckBoxOn = treeModel->mayaQtUtil().createPixmap(":/ImportDialog/checkboxOn.png");
-            fsCheckBoxOnDisabled
-                = treeModel->mayaQtUtil().createPixmap(":/ImportDialog/checkboxOnDisabled.png");
-            fsCheckBoxOff = treeModel->mayaQtUtil().createPixmap(":/ImportDialog/checkboxOff.png");
-            fsCheckBoxOffDisabled
-                = treeModel->mayaQtUtil().createPixmap(":/ImportDialog/checkboxOffDisabled.png");
-        } else {
-            // The tree model should never be null, but we can recover here if it is.
-            TF_RUNTIME_ERROR("Unexpected null tree model");
-            fsCheckBoxOn = new QPixmap(":/ImportDialog/checkboxOn.png");
-            fsCheckBoxOnDisabled = new QPixmap(":/ImportDialog/checkboxOnDisabled.png");
-            fsCheckBoxOff = new QPixmap(":/ImportDialog/checkboxOff.png");
-            fsCheckBoxOffDisabled = new QPixmap(":/ImportDialog/checkboxOffDisabled.png");
-        }
+        fsCheckBoxOn = createPixmap(":/ImportDialog/checkboxOn.png");
+        fsCheckBoxOnDisabled = createPixmap(":/ImportDialog/checkboxOnDisabled.png");
+        fsCheckBoxOff = createPixmap(":/ImportDialog/checkboxOff.png");
+        fsCheckBoxOffDisabled = createPixmap(":/ImportDialog/checkboxOffDisabled.png");
     }
 
     switch (fCheckState) {
@@ -76,30 +85,34 @@ const QPixmap& TreeItem::checkImage() const
 
 void TreeItem::setCheckState(TreeItem::CheckState st)
 {
-    assert(fType == Type::kLoad);
-    if (fType == Type::kLoad)
+    assert(fColumn == kColumnLoad);
+    if (fColumn == kColumnLoad)
         fCheckState = st;
 }
 
 void TreeItem::setVariantSelectionModified()
 {
-    assert(fType == Type::kVariants);
-    if (fType == Type::kVariants)
+    assert(fColumn == kColumnVariants);
+    if (fColumn == kColumnVariants)
         fVariantSelectionModified = true;
 }
 
-void TreeItem::initializeItem()
+void TreeItem::initializeItem(bool isDefaultPrim)
 {
-    switch (fType) {
-    case Type::kLoad: fCheckState = CheckState::kChecked_Disabled; break;
-    case Type::kName:
+    switch (fColumn) {
+    case kColumnLoad: fCheckState = CheckState::kChecked_Disabled; break;
+    case kColumnName:
         if (fPrim.IsPseudoRoot())
             setText("Root");
         else
             setText(QString::fromStdString(fPrim.GetName().GetString()));
+        if (isDefaultPrim) {
+            if (const QPixmap* pixmap = TreeModel::getDefaultPrimPixmap())
+                setData(*pixmap, Qt::DecorationRole);
+        }
         break;
-    case Type::kType: setText(QString::fromStdString(fPrim.GetTypeName().GetString())); break;
-    case Type::kVariants: {
+    case kColumnType: setText(QString::fromStdString(fPrim.GetTypeName().GetString())); break;
+    case kColumnVariants: {
         if (fPrim.HasVariantSets()) {
             // We set a special role flag when this prim has variant sets.
             // So we know when to create the label and combo box(es) for the variant