Change CanExport function to take MObject and updated python wrapper

Autodesk · Jul 11, 2023 · a43f3e8 · a43f3e8
1 parent ac4f5e0
commit a43f3e8
Show file tree

Hide file tree

Showing 14 changed files with 65 additions and 49 deletions.
diff --git a/lib/mayaUsd/fileio/primWriter.h b/lib/mayaUsd/fileio/primWriter.h
@@ -78,7 +78,7 @@ class UsdMayaPrimWriter
     /// declaring how well this class can support the current context.
     MAYAUSD_CORE_PUBLIC
     static ContextSupport
-    CanExport(const UsdMayaJobExportArgs& exportArgs, const MFnDependencyNode& exportNode);
+    CanExport(const UsdMayaJobExportArgs& exportArgs, const MObject& exportObj);
 
     MAYAUSD_CORE_PUBLIC
     virtual ~UsdMayaPrimWriter();

diff --git a/lib/mayaUsd/fileio/primWriterRegistry.cpp b/lib/mayaUsd/fileio/primWriterRegistry.cpp
@@ -59,15 +59,15 @@ static int                                                   _indexCounter = 0;
 _Registry::const_iterator _Find(
     const std::string&          mayaTypeName,
     const UsdMayaJobExportArgs& exportArgs,
-    const MFnDependencyNode&    exportNode)
+    const MObject&    exportObj)
 {
     using ContextSupport = UsdMayaPrimWriter::ContextSupport;
 
     _Registry::const_iterator ret = _reg.cend();
     _Registry::const_iterator first, last;
     std::tie(first, last) = _reg.equal_range(mayaTypeName);
     while (first != last) {
-        ContextSupport support = first->second._pred(exportArgs, exportNode);
+        ContextSupport support = first->second._pred(exportArgs, exportObj);
         // Look for a "Supported" reader. If no "Supported" reader is found, use a "Fallback" reader
         if (support == ContextSupport::Supported) {
             ret = first;
@@ -126,7 +126,7 @@ void UsdMayaPrimWriterRegistry::Register(
     // Use default ContextSupport if not specified
     _reg.insert(std::make_pair(
         mayaTypeName,
-        _RegistryEntry { [](const UsdMayaJobExportArgs&, const MFnDependencyNode&) {
+        _RegistryEntry { [](const UsdMayaJobExportArgs&, const MObject&) {
                             return UsdMayaPrimWriter::ContextSupport::Fallback;
                         },
                          fn,
@@ -160,11 +160,11 @@ void UsdMayaPrimWriterRegistry::RegisterRaw(
 UsdMayaPrimWriterRegistry::WriterFactoryFn UsdMayaPrimWriterRegistry::Find(
     const std::string&          mayaTypeName,
     const UsdMayaJobExportArgs& exportArgs,
-    const MFnDependencyNode&    exportNode)
+    const MObject&              exportObj)
 {
     TfRegistryManager::GetInstance().SubscribeTo<UsdMayaPrimWriterRegistry>();
 
-    _Registry::const_iterator it = _Find(mayaTypeName, exportArgs, exportNode);
+    _Registry::const_iterator it = _Find(mayaTypeName, exportArgs, exportObj);
 
     if (it != _reg.end()) {
         return it->second._writer;
@@ -173,7 +173,7 @@ UsdMayaPrimWriterRegistry::WriterFactoryFn UsdMayaPrimWriterRegistry::Find(
     static const TfTokenVector SCOPE = { _tokens->UsdMaya, _tokens->PrimWriter };
     UsdMaya_RegistryHelper::FindAndLoadMayaPlug(SCOPE, mayaTypeName);
 
-    it = _Find(mayaTypeName, exportArgs, exportNode);
+    it = _Find(mayaTypeName, exportArgs, exportObj);
 
     if (it != _reg.end()) {
         return it->second._writer;
@@ -185,7 +185,7 @@ UsdMayaPrimWriterRegistry::WriterFactoryFn UsdMayaPrimWriterRegistry::Find(
         // Nothing registered at all, remember that:
         _reg.insert(std::make_pair(
             mayaTypeName,
-            _RegistryEntry { [](const UsdMayaJobExportArgs&, const MFnDependencyNode&) {
+            _RegistryEntry { [](const UsdMayaJobExportArgs&, const MObject&) {
                                 return UsdMayaPrimWriter::ContextSupport::Fallback;
                             },
                              nullptr,

diff --git a/lib/mayaUsd/fileio/primWriterRegistry.h b/lib/mayaUsd/fileio/primWriterRegistry.h
@@ -80,7 +80,7 @@ struct UsdMayaPrimWriterRegistry
     /// Predicate function, i.e. a function that can tell the level of support
     /// the writer function will provide for a given set of export options.
     using ContextPredicateFn = std::function<
-        UsdMayaPrimWriter::ContextSupport(const UsdMayaJobExportArgs&, const MFnDependencyNode&)>;
+        UsdMayaPrimWriter::ContextSupport(const UsdMayaJobExportArgs&, const MObject&)>;
 
     /// \brief Register \p fn as a factory function providing a
     /// UsdMayaPrimWriter subclass that can be used to write \p mayaType.
@@ -145,7 +145,7 @@ struct UsdMayaPrimWriterRegistry
     static WriterFactoryFn Find(
         const std::string&          mayaTypeName,
         const UsdMayaJobExportArgs& exportArgs,
-        const MFnDependencyNode&    exportNode);
+        const MObject&    exportObj);
 
     /// \brief Check for external primWriter for \p mayaTypeName.
     MAYAUSD_CORE_PUBLIC

diff --git a/lib/mayaUsd/fileio/shaderWriterRegistry.cpp b/lib/mayaUsd/fileio/shaderWriterRegistry.cpp
@@ -57,7 +57,7 @@ static int       _indexCounter = 0;
 
 _Registry::const_iterator _Find(const TfToken& usdInfoId, const UsdMayaJobExportArgs& exportArgs)
 {
-    using ContextSupport = UsdMayaShaderWriter::ContextSupport;
+    using ContextSupport = UsdMayaPrimWriter::ContextSupport;
 
     TfToken    conversion = exportArgs.convertMaterialsTo;
     const bool noFallback
@@ -141,7 +141,7 @@ UsdMayaShaderWriterRegistry::WriterFactoryFn UsdMayaShaderWriterRegistry::Find(
         _reg.insert(std::make_pair(
             mayaTypeName,
             _RegistryEntry { [](const UsdMayaJobExportArgs&) {
-                                return UsdMayaShaderWriter::ContextSupport::Fallback;
+                                return UsdMayaPrimWriter::ContextSupport::Fallback;
                             },
                              nullptr,
                              -1 }));

diff --git a/lib/mayaUsd/fileio/shaderWriterRegistry.h b/lib/mayaUsd/fileio/shaderWriterRegistry.h
@@ -84,7 +84,7 @@ struct UsdMayaShaderWriterRegistry
     /// Predicate function, i.e. a function that can tell the level of support
     /// the writer function will provide for a given set of export options.
     using ContextPredicateFn
-        = std::function<UsdMayaShaderWriter::ContextSupport(const UsdMayaJobExportArgs&)>;
+        = std::function<UsdMayaPrimWriter::ContextSupport(const UsdMayaJobExportArgs&)>;
 
     /// \brief Register \p fn as a factory function providing a
     /// UsdMayaShaderWriter subclass that can be used to write \p mayaType.

diff --git a/lib/mayaUsd/fileio/shading/symmetricShaderWriter.cpp b/lib/mayaUsd/fileio/shading/symmetricShaderWriter.cpp
@@ -67,7 +67,7 @@ void UsdMayaSymmetricShaderWriter::RegisterWriter(
 }
 
 /* static */
-UsdMayaShaderWriter::ContextSupport UsdMayaSymmetricShaderWriter::CanExport(
+UsdMayaPrimWriter::ContextSupport UsdMayaSymmetricShaderWriter::CanExport(
     const UsdMayaJobExportArgs& exportArgs,
     const TfToken&              materialConversionName)
 {

diff --git a/lib/mayaUsd/fileio/writeJobContext.cpp b/lib/mayaUsd/fileio/writeJobContext.cpp
@@ -565,7 +565,7 @@ UsdMayaWriteJobContext::_FindWriter(const MFnDependencyNode& mayaNode)
         = UsdMayaUtil::GetAllAncestorMayaNodeTypes(mayaNodeType);
     for (auto i = ancestorTypes.rbegin(); i != ancestorTypes.rend(); ++i) {
         if (UsdMayaPrimWriterRegistry::WriterFactoryFn primWriterFactory
-            = UsdMayaPrimWriterRegistry::Find(*i, mArgs, mayaNode)) {
+            = UsdMayaPrimWriterRegistry::Find(*i, mArgs, mayaNode.object())) {
             mWriterFactoryCache[mayaNodeType] = primWriterFactory;
             return primWriterFactory;
         }

diff --git a/lib/mayaUsd/python/wrapPrimWriter.cpp b/lib/mayaUsd/python/wrapPrimWriter.cpp
@@ -206,20 +206,32 @@ class PrimWriterWrapper
             return sptr;
         }
 
+        UsdMayaPrimWriter::ContextSupport
+        operator()(const UsdMayaJobExportArgs& exportArgs, const MObject& exportObj)
+        {
+
+            boost::python::object pyClass = GetPythonObject(_classIndex);
+            if (!pyClass) {
+                // Prototype was unregistered
+                return UsdMayaPrimWriter::ContextSupport::Unsupported;
+            }
+            TfPyLock              pyLock;
+            boost::python::object CanExport = pyClass.attr("CanExport");
+            PyObject*             callable = CanExport.ptr();
+            auto                  res = boost::python::call<int>(callable, exportArgs, exportObj);
+            return UsdMayaPrimWriter::ContextSupport(res);            
+        }
+
         // Create a new wrapper for a Python class that is seen for the first time for a given
         // purpose. If we already have a registration for this purpose: update the class to
         // allow the previously issued factory function to use it.
-        static UsdMayaPrimWriterRegistry::WriterFactoryFn
-        Register(boost::python::object cl, const std::string& mayaTypeName)
+        static FactoryFnWrapper
+        Register(boost::python::object cl, const std::string& mayaTypeName, bool& updated)
         {
             size_t classIndex = RegisterPythonObject(cl, GetKey(cl, mayaTypeName));
-            if (classIndex != UsdMayaPythonObjectRegistry::UPDATED) {
-                // Return a new factory function:
-                return FactoryFnWrapper { classIndex };
-            } else {
-                // We already registered a factory function for this purpose:
-                return nullptr;
-            }
+            updated = classIndex == UsdMayaPythonObjectRegistry::UPDATED;
+            // Return a new factory function:
+            return FactoryFnWrapper { classIndex };
         }
 
         // Unregister a class for a given purpose. This will cause the associated factory
@@ -246,10 +258,10 @@ class PrimWriterWrapper
 
     static void Register(boost::python::object cl, const std::string& mayaTypeName)
     {
-        UsdMayaPrimWriterRegistry::WriterFactoryFn fn
-            = FactoryFnWrapper::Register(cl, mayaTypeName);
-        if (fn) {
-            UsdMayaPrimWriterRegistry::Register(mayaTypeName, fn, true);
+        bool             updated = false;
+        FactoryFnWrapper fn = FactoryFnWrapper::Register(cl, mayaTypeName, updated);
+        if (!updated) {
+            UsdMayaPrimWriterRegistry::Register(mayaTypeName, fn, fn, true);
         }
     }
 
@@ -349,18 +361,18 @@ class ShaderWriterWrapper : public PrimWriterWrapper<UsdMayaShaderWriter>
         }
 
         // We can have multiple function objects, this one apapts the CanImport function:
-        UsdMayaShaderWriter::ContextSupport operator()(const UsdMayaJobExportArgs& exportArgs)
+        UsdMayaPrimWriter::ContextSupport operator()(const UsdMayaJobExportArgs& exportArgs)
         {
             boost::python::object pyClass = GetPythonObject(_classIndex);
             if (!pyClass) {
                 // Prototype was unregistered
-                return UsdMayaShaderWriter::ContextSupport::Unsupported;
+                return UsdMayaPrimWriter::ContextSupport::Unsupported;
             }
             TfPyLock              pyLock;
             boost::python::object CanExport = pyClass.attr("CanExport");
             PyObject*             callable = CanExport.ptr();
             auto                  res = boost::python::call<int>(callable, exportArgs);
-            return UsdMayaShaderWriter::ContextSupport(res);
+            return UsdMayaPrimWriter::ContextSupport(res);
         }
 
         // Create a new wrapper for a Python class that is seen for the first time for a given
@@ -598,11 +610,24 @@ void wrapJobExportArgs()
         .staticmethod("GetDefaultMaterialsScopeName");
 }
 
+TF_REGISTRY_FUNCTION(TfEnum)
+{
+    TF_ADD_ENUM_NAME(UsdMayaPrimWriter::ContextSupport::Supported, "Supported");
+    TF_ADD_ENUM_NAME(UsdMayaPrimWriter::ContextSupport::Fallback, "Fallback");
+    TF_ADD_ENUM_NAME(UsdMayaPrimWriter::ContextSupport::Unsupported, "Unsupported");
+}
+
 void wrapPrimWriter()
 {
-    boost::python::class_<PrimWriterWrapper<>, boost::noncopyable>(
-        "PrimWriter", boost::python::no_init)
-        .def("__init__", make_constructor(&PrimWriterWrapper<>::New))
+
+    boost::python::class_<PrimWriterWrapper<>, boost::noncopyable> c(
+        "PrimWriter", boost::python::no_init);
+
+    boost::python::scope s(c);
+
+    TfPyWrapEnum<UsdMayaPrimWriter::ContextSupport>();
+
+    c.def("__init__", make_constructor(&PrimWriterWrapper<>::New))
         .def(
             "PostExport",
             &PrimWriterWrapper<>::PostExport,
@@ -658,13 +683,6 @@ void wrapPrimWriter()
         .staticmethod("Unregister");
 }
 
-TF_REGISTRY_FUNCTION(TfEnum)
-{
-    TF_ADD_ENUM_NAME(UsdMayaShaderWriter::ContextSupport::Supported, "Supported");
-    TF_ADD_ENUM_NAME(UsdMayaShaderWriter::ContextSupport::Fallback, "Fallback");
-    TF_ADD_ENUM_NAME(UsdMayaShaderWriter::ContextSupport::Unsupported, "Unsupported");
-}
-
 //----------------------------------------------------------------------------------------------------------------------
 void wrapShaderWriter()
 {
@@ -674,8 +692,6 @@ void wrapShaderWriter()
 
     boost::python::scope s(c);
 
-    TfPyWrapEnum<UsdMayaShaderWriter::ContextSupport>();
-
     c.def("__init__", make_constructor(&ShaderWriterWrapper::New))
         .def(
             "GetShadingAttributeNameForMayaAttrName",

diff --git a/lib/usd/pxrUsdPreviewSurface/usdPreviewSurfaceWriter.cpp b/lib/usd/pxrUsdPreviewSurface/usdPreviewSurfaceWriter.cpp
@@ -135,7 +135,7 @@ void _ValidateNormalMap(const MPlug& shadingNodePlug)
 
 } // namespace
 
-UsdMayaShaderWriter::ContextSupport
+UsdMayaPrimWriter::ContextSupport
 PxrMayaUsdPreviewSurface_Writer::CanExport(const UsdMayaJobExportArgs& exportArgs)
 {
     if (exportArgs.convertMaterialsTo == UsdImagingTokens->UsdPreviewSurface) {

diff --git a/lib/usd/translators/shading/mtlxBaseWriter.cpp b/lib/usd/translators/shading/mtlxBaseWriter.cpp
@@ -75,7 +75,7 @@ REGISTER_SHADING_MODE_EXPORT_MATERIAL_CONVERSION(
     TrMtlxTokens->niceName,
     TrMtlxTokens->exportDescription);
 
-UsdMayaShaderWriter::ContextSupport
+UsdMayaPrimWriter::ContextSupport
 MtlxUsd_BaseWriter::CanExport(const UsdMayaJobExportArgs& exportArgs)
 {
     return exportArgs.convertMaterialsTo == TrMtlxTokens->conversionName

diff --git a/lib/usd/translators/shading/mtlxSymmetricShaderWriter.cpp b/lib/usd/translators/shading/mtlxSymmetricShaderWriter.cpp
@@ -146,7 +146,7 @@ void MtlxUsd_SymmetricShaderWriter::RegisterWriter(
 }
 
 /* static */
-UsdMayaShaderWriter::ContextSupport
+UsdMayaPrimWriter::ContextSupport
 MtlxUsd_SymmetricShaderWriter::CanExport(const UsdMayaJobExportArgs& exportArgs)
 {
     if (exportArgs.convertMaterialsTo == TrMtlxTokens->conversionName) {

diff --git a/lib/usd/translators/shading/usdDisplacementShaderWriter.cpp b/lib/usd/translators/shading/usdDisplacementShaderWriter.cpp
@@ -152,7 +152,7 @@ void PxrUsdTranslators_DisplacementShaderWriter::Write(const UsdTimeCode& usdTim
 }
 
 /* static */
-UsdMayaShaderWriter::ContextSupport
+UsdMayaPrimWriter::ContextSupport
 PxrUsdTranslators_DisplacementShaderWriter::CanExport(const UsdMayaJobExportArgs& exportArgs)
 {
     if (exportArgs.convertMaterialsTo == UsdImagingTokens->UsdPreviewSurface) {

diff --git a/lib/usd/translators/shading/usdFileTextureWriter.cpp b/lib/usd/translators/shading/usdFileTextureWriter.cpp
@@ -115,7 +115,7 @@ TF_DEFINE_PRIVATE_TOKENS(
 );
 // clang-format on
 
-UsdMayaShaderWriter::ContextSupport
+UsdMayaPrimWriter::ContextSupport
 PxrUsdTranslators_FileTextureWriter::CanExport(const UsdMayaJobExportArgs& exportArgs)
 {
     if (exportArgs.convertMaterialsTo == UsdImagingTokens->UsdPreviewSurface) {

diff --git a/lib/usd/translators/shading/usdMaterialWriter.cpp b/lib/usd/translators/shading/usdMaterialWriter.cpp
@@ -51,7 +51,7 @@
 
 PXR_NAMESPACE_OPEN_SCOPE
 
-UsdMayaShaderWriter::ContextSupport
+UsdMayaPrimWriter::ContextSupport
 PxrUsdTranslators_MaterialWriter::CanExport(const UsdMayaJobExportArgs& exportArgs)
 {
     if (exportArgs.convertMaterialsTo == UsdImagingTokens->UsdPreviewSurface) {