Module classes are not forced to be called "Module"

Signed-off-by: bhavishyagopesh <bhavishyagopesh@gmail.com>

src/mgr/MgrPyModule.cc

@@ -171,30 +171,45 @@ int MgrPyModule::load()
   }
 
   // Find the class
-  // TODO: let them call it what they want instead of just 'Module'
-  auto pClass = PyObject_GetAttrString(pModule, (const char*)"Module");
-  Py_DECREF(pModule);
-  if (pClass == nullptr) {
+  auto  mgr_module_type = PyObject_GetAttrString(pModule, (const char*)"MgrModule");
+
+  if (mgr_module_type == nullptr) {
     derr << "Class not found in module '" << module_name << "'" << dendl;
     derr << handle_pyerror() << dendl;
     return -EINVAL;
   }
 
-
-  // Just using the module name as the handle, replace with a
-  // uuidish thing if needed
-  auto pyHandle = PyString_FromString(module_name.c_str());
-  auto pArgs = PyTuple_Pack(1, pyHandle);
-  pClassInstance = PyObject_CallObject(pClass, pArgs);
-  Py_DECREF(pClass);
-  Py_DECREF(pyHandle);
-  Py_DECREF(pArgs);
-  if (pClassInstance == nullptr) {
-    derr << "Failed to construct class in '" << module_name << "'" << dendl;
-    derr << handle_pyerror() << dendl;
-    return -EINVAL;
-  } else {
-    dout(1) << "Constructed class from module: " << module_name << dendl;
+  Py_DECREF(pModule);
+
+  auto locals = PyModule_GetDict(pModule);
+  PyObject *key, *value;
+  Py_ssize_t pos = 0;
+
+  while (PyDict_Next(locals, &pos, &key, &value)) {
+
+    if (!PyType_Check(value)) {
+      continue;
+    }
+
+    if (PyObject_IsSubclass(value, mgr_module_type)) {
+
+        // Just using the module name as the handle, replace with a
+        // uuidish thing if needed
+        auto pClass = value;
+        auto pyHandle = PyString_FromString(module_name.c_str());
+        auto pArgs = PyTuple_Pack(1, pyHandle);
+        pClassInstance = PyObject_CallObject(pClass, pArgs);
+        Py_DECREF(pyHandle);
+        Py_DECREF(pArgs);
+        if (pClassInstance == nullptr) {
+          derr << "Failed to construct class '" << PyObject_Str(pClass) << "'" << " in " << "'" << module_name << "'" << dendl;
+          derr << handle_pyerror() << dendl;
+          return -EINVAL;
+        } else {
+          dout(1) << "Constructed class: " << PyObject_Str(pClass) << " from module: " << module_name << dendl;
+        }
+        Py_DECREF(pClass);
+    }
   }
 
   return load_commands();