Merge pull request #2044 from Autodesk/t_gamaj/MAYA-113541/allow_regi…

…stration_updates

MAYA-113541 Allow updating registered Python class

lib/mayaUsd/python/CMakeLists.txt

@@ -25,6 +25,7 @@ add_library(${PYTHON_TARGET_NAME} SHARED)
 target_sources(${PYTHON_TARGET_NAME} 
     PRIVATE
         module.cpp
+        pythonObjectRegistry.cpp
         wrapSparseValueWriter.cpp
         wrapAdaptor.cpp
         wrapBlockSceneModificationContext.cpp

lib/mayaUsd/python/pythonObjectRegistry.cpp

@@ -0,0 +1,107 @@
+//
+// Copyright 2022 Autodesk
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include "pythonObjectRegistry.h"
+
+#include <pxr/base/tf/diagnostic.h>
+#include <pxr/base/tf/pyUtils.h>
+#include <pxr/pxr.h>
+
+#include <boost/python.hpp>
+#include <boost/python/def.hpp>
+
+using namespace std;
+using namespace boost::python;
+using namespace boost;
+
+PXR_NAMESPACE_USING_DIRECTIVE
+
+// Registers or updates a Python class for the provided key.
+size_t UsdMayaPythonObjectRegistry::RegisterPythonObject(object cl, const std::string& key)
+{
+    TClassIndex::const_iterator target = _sIndex.find(key);
+    if (target == _sIndex.cend()) {
+        HookInterpreterExit();
+        // Create a new entry:
+        size_t classIndex = _sClassVec.size();
+        _sClassVec.emplace_back(cl);
+        _sIndex.emplace(key, classIndex);
+        // Return a new factory function:
+        return classIndex;
+    } else {
+        _sClassVec[target->second] = cl;
+        return UPDATED;
+    }
+}
+
+void UsdMayaPythonObjectRegistry::UnregisterPythonObject(object cl, const std::string& key)
+{
+    TClassIndex::const_iterator target = _sIndex.find(key);
+    if (target != _sIndex.cend()) {
+        // Clear the Python class:
+        _sClassVec[target->second] = object();
+    }
+}
+
+bool UsdMayaPythonObjectRegistry::IsPythonClass(object cl)
+{
+    auto classAttr = cl.attr("__class__");
+    if (!classAttr) {
+        return false;
+    }
+    auto nameAttr = classAttr.attr("__name__");
+    if (!nameAttr) {
+        return false;
+    }
+    std::string name = extract<std::string>(nameAttr);
+    return name == "class";
+}
+
+std::string UsdMayaPythonObjectRegistry::ClassName(object cl)
+{
+    // Is it a Python class:
+    if (!IsPythonClass(cl)) {
+        // So far class is always first parameter, so we can have this check be here.
+        TfPyThrowRuntimeError("First argument must be a Python class");
+    }
+
+    auto nameAttr = cl.attr("__name__");
+    if (!nameAttr) {
+        TfPyThrowRuntimeError("Unexpected Python error: No __name__ attribute");
+    }
+
+    return std::string(boost::python::extract<std::string>(nameAttr));
+}
+
+UsdMayaPythonObjectRegistry::TClassVec   UsdMayaPythonObjectRegistry::_sClassVec;
+UsdMayaPythonObjectRegistry::TClassIndex UsdMayaPythonObjectRegistry::_sIndex;
+
+void UsdMayaPythonObjectRegistry::HookInterpreterExit()
+{
+    if (_sClassVec.empty()) {
+        if (import("atexit").attr("register")(&OnInterpreterExit).is_none()) {
+            TF_CODING_ERROR("Couldn't register unloader to atexit");
+        }
+    }
+}
+
+void UsdMayaPythonObjectRegistry::OnInterpreterExit()
+{
+    // Release all Python classes:
+    for (size_t i = 0; i < _sClassVec.size(); ++i) {
+        _sClassVec[i] = object();
+    }
+}

lib/mayaUsd/python/pythonObjectRegistry.h

@@ -0,0 +1,66 @@
+//
+// Copyright 2022 Autodesk
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+#ifndef USDMAYA_PYTHON_CLASS_REGISTRY_H
+#define USDMAYA_PYTHON_CLASS_REGISTRY_H
+
+#include <pxr/pxr.h>
+
+#include <boost/python/class.hpp>
+
+#include <map>
+#include <string>
+#include <vector>
+
+PXR_NAMESPACE_OPEN_SCOPE
+
+//---------------------------------------------------------------------------------------------
+/// \brief Keeps track of registered Python classes and allows updating and unregistering them
+//---------------------------------------------------------------------------------------------
+class UsdMayaPythonObjectRegistry
+{
+protected:
+    static const size_t UPDATED = 0xFFFFFFFF;
+
+    // Registers or updates a Python class for the provided key.
+    static size_t RegisterPythonObject(boost::python::object cl, const std::string& key);
+
+    // Unregister a Python class for a given key. This will cause the associated factory function to
+    // stop producing this Python class.
+    static void UnregisterPythonObject(boost::python::object cl, const std::string& key);
+
+    static boost::python::object GetPythonObject(size_t index) { return _sClassVec[index]; }
+
+    static bool IsPythonClass(boost::python::object cl);
+
+    static std::string ClassName(boost::python::object cl);
+
+private:
+    // Static table of all registered Python classes, with associated index:
+    typedef std::vector<boost::python::object> TClassVec;
+    static TClassVec                           _sClassVec;
+    typedef std::map<std::string, size_t>      TClassIndex;
+    static TClassIndex                         _sIndex;
+
+    static void HookInterpreterExit();
+
+public:
+    // To be called when the Python interpreter exits.
+    static void OnInterpreterExit();
+};
+
+PXR_NAMESPACE_CLOSE_SCOPE
+
+#endif

lib/mayaUsd/python/wrapExportChaser.cpp

@@ -14,6 +14,8 @@
 // limitations under the License.
 //
 
+#include "pythonObjectRegistry.h"
+
 #include <mayaUsd/fileio/chaser/exportChaser.h>
 #include <mayaUsd/fileio/chaser/exportChaserRegistry.h>
 #include <mayaUsd/fileio/registryHelper.h>
@@ -65,19 +67,80 @@ class ExportChaserWrapper
         return this->CallVirtual<>("PostExport", &This::default_PostExport)();
     }
 
+    //---------------------------------------------------------------------------------------------
+    /// \brief  wraps a factory function that allows registering an updated Python class
+    //---------------------------------------------------------------------------------------------
+    class FactoryFnWrapper : public UsdMayaPythonObjectRegistry
+    {
+    public:
+        // Instances of this class act as "function objects" that are fully compatible with the
+        // std::function requested by UsdMayaSchemaApiAdaptorRegistry::Register. These will create
+        // python wrappers based on the latest Class registered.
+        UsdMayaExportChaser*
+        operator()(const UsdMayaExportChaserRegistry::FactoryContext& factoryContext)
+        {
+            boost::python::object pyClass = GetPythonObject(_classIndex);
+            if (!pyClass) {
+                // Prototype was unregistered
+                return nullptr;
+            }
+            auto                  chaser = new ExportChaserWrapper();
+            TfPyLock              pyLock;
+            boost::python::object instance = pyClass(factoryContext, (uintptr_t)chaser);
+            boost::python::incref(instance.ptr());
+            initialize_wrapper(instance.ptr(), chaser);
+            return chaser;
+        }
+
+        // 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 UsdMayaExportChaserRegistry::FactoryFn
+        Register(boost::python::object cl, const std::string& mayaTypeName)
+        {
+            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;
+            }
+        }
+
+        // Unregister a class for a given purpose. This will cause the associated factory
+        // function to stop producing this Python class.
+        static void Unregister(boost::python::object cl, const std::string& mayaTypeName)
+        {
+            UnregisterPythonObject(cl, GetKey(cl, mayaTypeName));
+        }
+
+    private:
+        // Function object constructor. Requires only the index of the Python class to use.
+        FactoryFnWrapper(size_t classIndex)
+            : _classIndex(classIndex) {};
+
+        size_t _classIndex;
+
+        // Generates a unique key based on the name of the class, along with the class
+        // purpose:
+        static std::string GetKey(boost::python::object cl, const std::string& mayaTypeName)
+        {
+            return ClassName(cl) + "," + mayaTypeName + "," + ",ExportChaser";
+        }
+    };
+
     static void Register(boost::python::object cl, const std::string& mayaTypeName)
     {
-        UsdMayaExportChaserRegistry::GetInstance().RegisterFactory(
-            mayaTypeName,
-            [=](const UsdMayaExportChaserRegistry::FactoryContext& factoryContext) {
-                auto                  chaser = new ExportChaserWrapper();
-                TfPyLock              pyLock;
-                boost::python::object instance = cl(factoryContext, (uintptr_t)chaser);
-                boost::python::incref(instance.ptr());
-                initialize_wrapper(instance.ptr(), chaser);
-                return chaser;
-            },
-            true);
+        UsdMayaExportChaserRegistry::FactoryFn fn = FactoryFnWrapper::Register(cl, mayaTypeName);
+        if (fn) {
+            UsdMayaExportChaserRegistry::GetInstance().RegisterFactory(mayaTypeName, fn, true);
+        }
+    }
+
+    static void Unregister(boost::python::object cl, const std::string& mayaTypeName)
+    {
+        FactoryFnWrapper::Unregister(cl, mayaTypeName);
     }
 };
 
@@ -109,5 +172,7 @@ void wrapExportChaser()
         .def("ExportFrame", &This::ExportFrame, &ExportChaserWrapper::default_ExportFrame)
         .def("PostExport", &This::PostExport, &ExportChaserWrapper::default_PostExport)
         .def("Register", &ExportChaserWrapper::Register)
-        .staticmethod("Register");
+        .staticmethod("Register")
+        .def("Unregister", &ExportChaserWrapper::Unregister)
+        .staticmethod("Unregister");
 }