Implement loading directly into a Maya USD Proxy Shape

client/ayon_maya/api/lib.py

@@ -1864,6 +1864,10 @@ def get_container_members(container, include_reference_associated_nodes=False):
         # Assume it's a container dictionary
         container = container["objectName"]
 
+    if "," in container:
+        # Assume it's a UFE path - return it as the only member
+        return [container]
+
     members = cmds.sets(container, query=True) or []
     members = cmds.ls(members, long=True, objectsOnly=True) or []
     all_members = set(members)

client/ayon_maya/api/pipeline.py

@@ -408,6 +408,36 @@ def parse_container(container):
     return data
 
 
+def parse_usd_prim_container(prim, proxy):
+    """Parse instance container from UsdPrim if it is marked as one
+
+    Args:
+        prim (pxr.Usd.Prim): USD Primitive
+        proxy (str): The maya usd stage proxy shape node the primitive
+            belongs to.
+
+    Returns:
+        dict: The container schema data for this container node.
+
+    """
+    data = prim.GetCustomDataByKey("ayon")
+    if not data or not data.get("id") in {AVALON_CONTAINER_ID,
+                                          AYON_CONTAINER_ID}:
+        return
+
+    # Store transient data
+    data["prim"] = prim
+    data["proxy"] = proxy
+
+    # Store the maya UFE path as objectName
+    prim_path = str(prim.GetPath())
+    data["objectName"] = "{},{}".format(proxy, prim_path)
+    data["namespace"] = prim_path
+    data["name"] = proxy
+
+    return data
+
+
 def _ls():
     """Yields AYON container node names.
 
@@ -449,6 +479,28 @@ def _maya_iterate(iterator):
         if value in ids:
             yield fn_dep.name()
 
+    for container in ls_maya_usd_proxy_prims():
+        yield container
+
+
+def ls_maya_usd_proxy_prims():
+    # TODO: This might be nicer once the Loader API gets a refactor where
+    #  the loaders themselves can return the containers from the scene
+    if cmds.pluginInfo("mayaUsdPlugin", query=True, loaded=True):
+        usd_proxies = cmds.ls(type="mayaUsdProxyShape", long=True)
+        if usd_proxies:
+            import mayaUsd.ufe
+            for proxy in usd_proxies:
+                stage = mayaUsd.ufe.getStage(proxy)
+                if not stage:
+                    log.warning(f"Unable to get USD stage from: {proxy}")
+                    continue
+
+                for prim in stage.TraverseAll():
+                    container = parse_usd_prim_container(prim, proxy=proxy)
+                    if container:
+                        yield container
+
 
 def ls():
     """Yields containers from active Maya scene

client/ayon_maya/api/usdlib.py

@@ -0,0 +1,80 @@
+from ayon_core.pipeline.constants import AVALON_CONTAINER_ID
+from maya import cmds
+from pxr import Sdf
+
+
+def remove_spec(spec):
+    """Delete Sdf.PrimSpec or Sdf.PropertySpec
+
+    Also see:
+        https://forum.aousd.org/t/api-basics-for-designing-a-manage-edits-editor-for-usd/676/1  # noqa
+        https://gist.github.com/BigRoy/4d2bf2eef6c6a83f4fda3c58db1489a5
+
+    """
+    if spec.expired:
+        return
+
+    if isinstance(spec, Sdf.PrimSpec):
+        # PrimSpec
+        parent = spec.nameParent
+        if parent:
+            view = parent.nameChildren
+        else:
+            # Assume PrimSpec is root prim
+            view = spec.layer.rootPrims
+        del view[spec.name]
+
+    elif isinstance(spec, Sdf.PropertySpec):
+        # Relationship and Attribute specs
+        del spec.owner.properties[spec.name]
+    else:
+        raise TypeError(f"Unsupported spec type: {spec}")
+
+
+def iter_ufe_usd_selection():
+    """Yield Maya USD Proxy Shape related UFE paths in selection.
+
+    The returned path are the Maya node name joined by a command to the
+    USD prim path.
+
+    Yields:
+        str: Path to UFE path in USD stage in selection.
+
+    """
+    for path in cmds.ls(selection=True, ufeObjects=True, long=True,
+                        absoluteName=True):
+        if "," not in path:
+            continue
+
+        node, ufe_path = path.split(",", 1)
+        if cmds.nodeType(node) != "mayaUsdProxyShape":
+            continue
+
+        yield path
+
+
+def containerise_prim(prim,
+                      name,
+                      namespace,
+                      context,
+                      loader):
+    """Containerise a USD prim.
+
+    Arguments:
+        prim (pxr.Usd.Prim): The prim to containerise.
+        name (str): Name to containerize.
+        namespace (str): Namespace to containerize.
+        context (dict): Load context (incl. representation).
+        name (str): Name to containerize.
+        loader (str): Loader name.
+
+    """
+    for key, value in {
+        "ayon:schema": "openpype:container-2.0",
+        "ayon:id": AVALON_CONTAINER_ID,
+        "ayon:name": name,
+        "ayon:namespace": namespace,
+        "ayon:loader": loader,
+        "ayon:representation": context["representation"]["id"],
+    }.items():
+        prim.SetCustomDataByKey(key, str(value))

client/ayon_maya/plugins/load/load_maya_usd_add_maya_reference.py

@@ -0,0 +1,156 @@
+# -*- coding: utf-8 -*-
+import contextlib
+
+from ayon_core.pipeline import load
+from ayon_maya.api.usdlib import (
+    containerise_prim,
+    iter_ufe_usd_selection
+)
+
+from maya import cmds
+import mayaUsd
+
+
+@contextlib.contextmanager
+def no_edit_mode(prim, restore_after=True):
+    """Ensure MayaReference prim is not in edit mode during context"""
+    pulled_node = mayaUsd.lib.PrimUpdaterManager.readPullInformation(prim)
+    ufe_path = None
+    try:
+        # remove edit state if pulled
+        if pulled_node:
+            import mayaUsdUtils
+            assert mayaUsdUtils.isPulledMayaReference(pulled_node)
+            cmds.mayaUsdDiscardEdits(pulled_node)
+
+            # Discarding the edits directly selects the prim
+            # so we can get the UFE path from selection
+            ufe_path = cmds.ls(selection=True, ufeObjects=True, long=True)[0]
+
+        yield prim, ufe_path, pulled_node
+    finally:
+        if restore_after and pulled_node and ufe_path:
+            cmds.mayaUsdEditAsMaya(ufe_path)
+
+
+class MayaUsdProxyAddMayaReferenceLoader(load.LoaderPlugin):
+    """Read USD data in a Maya USD Proxy
+
+    TODO: It'd be much easier if this loader would be capable of returning the
+        available containers in the scene based on the AYON URLs inside a USD
+        stage. That way we could potentially avoid the need for custom metadata
+        keys, stay closer to USD native data and rely solely on the
+        AYON:asset=blue,subset=modelMain,version=1 url
+
+    """
+
+    product_types = {"*"}
+    representations = ["*"]
+    extensions = ["ma", "mb"]
+
+    label = "USD Add Maya Reference"
+    order = 1
+    icon = "code-fork"
+    color = "orange"
+
+    identifier_key = "ayon_identifier"
+
+    def load(self, context, name=None, namespace=None, options=None):
+
+        selection = list(iter_ufe_usd_selection())
+        assert len(selection) == 1, "Select only one PRIM please"
+        ufe_path = selection[0]
+        path = self.filepath_from_context(context)
+         # Make sure we can load the plugin
+        cmds.loadPlugin("mayaUsdPlugin", quiet=True)
+        import mayaUsdAddMayaReference
+
+        namespace = "test"
+        prim = mayaUsdAddMayaReference.createMayaReferencePrim(
+            ufe_path,
+            path,
+            namespace,
+            # todo: add more of the arguments
+            # mayaReferencePrimName Nameprim_name,
+            # groupPrim (3-tuple, group name, type and kind)
+            # variantSet (2-tuple, variant set name and variant name)
+        )
+        if not prim:
+            # Failed to add a reference
+            raise RuntimeError(f"Failed to add a reference at {ufe_path}")
+
+        containerise_prim(
+            prim,
+            name=name,
+            namespace=namespace or "",
+            context=context,
+            loader=self.__class__.__name__
+        )
+
+        return prim
+
+    def _update_reference_path(self, prim, filepath):
+        """Update MayaReference prim 'mayaReference' in nearest prim spec"""
+
+        from pxr import Sdf
+
+        # We want to update the authored opinion in the right place, e.g.
+        # within a VariantSet if it's authored there. We go through the
+        # PrimStack to find the first prim spec that authors an opinion
+        # on the 'mayaReference' attribute where we have permission to
+        # change it. This could technically mean we're altering it in
+        # layers that we might not want to (e.g. a published USD file?)
+        stack = prim.GetPrimStack()
+        for prim_spec in stack:
+            if "mayaReference" not in prim_spec.attributes:
+                # prim spec defines no opinion on mayaRefernce attribute?
+                continue
+
+            attr = prim_spec.attributes["mayaReference"]
+            if attr.permission != Sdf.PermissionPublic:
+                print(f"Not allowed to edit: {attr}")
+                continue
+
+            if filepath != attr.default:
+                print(
+                    f"Updating {attr.path} - {attr.default} -> {filepath}")
+                attr.default = filepath
+
+            # Attribute is either updated or already set to
+            # the value in that layer
+            return
+
+        # Just define in the current edit layer?
+        attr = prim.GetAttribute("mayaReference")
+        attr.Set(filepath)
+
+    def update(self, container, context):
+        # type: (dict, dict) -> None
+        """Update container with specified representation."""
+
+        prim = container["prim"]
+        representation = context["representation"]
+        filepath = self.filepath_from_context(context)
+
+        with no_edit_mode(prim):
+            self._update_reference_path(prim, filepath)
+
+        # Update representation id
+        # TODO: Do this in prim spec where we update reference path?
+        prim.SetCustomDataByKey(
+            "ayon:representation", str(representation["_id"])
+        )
+
+    def switch(self, container, context):
+        self.update(container, context)
+
+    def remove(self, container):
+        # type: (dict) -> None
+        """Remove loaded container."""
+
+        from ayon_maya.api.usdlib import remove_spec
+
+        prim = container["prim"]
+        with no_edit_mode(prim, restore_after=False):
+            for spec in prim.GetPrimStack():
+                remove_spec(spec)