Move GetRawAlignment to ExecutivePython

schrodinger · Jan 30, 2024 · 9e2ca7d · 9e2ca7d
1 parent ce1fb16
commit 9e2ca7d
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 66 deletions.
diff --git a/layer3/ExecutivePython.cpp b/layer3/ExecutivePython.cpp
@@ -256,4 +256,66 @@ pymol::Result<float> ExecutiveFitPairs(
   return ExecutiveRMSPairs(G, word, 2, quiet);
 }
 
+/**
+ * Implementation of ExecutiveGetRawAlignment
+ * @param alnobj alignment object
+ * @param active_only only consider active alignments
+ * @param state state of alignment object
+ * @return a list of lists of (object, index) tuples containing the
+ * raw per-atom alignment relationships
+ */
+static pymol::Result<PyObject*> ExecutiveGetRawAlignmentImpl(
+    PyMOLGlobals* G, const ObjectAlignment* alnobj, bool active_only, int state)
+{
+  if (state >= alnobj->getNFrame()) {
+    return pymol::make_error(
+        "Index Error: state ", state, " >= NState ", alnobj->getNFrame());
+  }
+
+  const auto& vla = alnobj->State[state].alignVLA;
+
+  if (!vla) {
+    return pymol::make_error("state ", state, " not valid");
+  }
+
+  auto hide_underscore = SettingGet<bool>(G, cSetting_hide_underscore_names);
+  const auto vla_len = VLAGetSize(vla);
+
+  PyObject* raw = PyList_New(0);
+
+  for (size_t i = 0; i < vla_len; ++i) {
+    PyObject* col = PyList_New(0);
+
+    for (int id; (id = vla[i]); ++i) {
+      auto eoo = ExecutiveUniqueIDAtomDictGet(G, id);
+      if (eoo && (!active_only || eoo->obj->Enabled) &&
+          (!hide_underscore || eoo->obj->Name[0] != '_')) {
+        PyObject* idx = Py_BuildValue("si", eoo->obj->Name, eoo->atm + 1);
+        PyList_Append(col, idx);
+        Py_DECREF(idx);
+      }
+    }
+
+    if (PyList_Size(col) > 0) {
+      PyList_Append(raw, col);
+    }
+    Py_DECREF(col);
+  }
+  return raw;
+}
+
+pymol::Result<PyObject*> ExecutiveGetRawAlignment(PyMOLGlobals* G,
+    pymol::null_safe_zstring_view name, bool active_only, int state)
+{
+  if (name.empty()) {
+    name = ExecutiveGetActiveAlignment(G);
+  }
+  if (!name.empty()) {
+    if (auto obj = ExecutiveFindObject<ObjectAlignment>(G, name)) {
+      return ExecutiveGetRawAlignmentImpl(G, obj, active_only, state);
+    }
+  }
+  return pymol::make_error("No such alignment: ", name);
+}
+
 #endif
diff --git a/layer3/ExecutivePython.h b/layer3/ExecutivePython.h
@@ -1,10 +1,12 @@
 #pragma once
 
+#ifndef _PYMOL_NO_PY
+
 #include "os_python.h"
 #include "Result.h"
 #include "pymol/zstring_view.h"
 
-#ifndef _PYMOL_NO_PY
+struct ObjectAlignment;
 
 pymol::Result<> ExecutiveLoadObject(PyMOLGlobals* G,
     const char* oname, PyObject* model, int frame, int type, int finish,
@@ -17,4 +19,14 @@ pymol::Result<> ExecutiveSetRawAlignment(PyMOLGlobals* G,
 pymol::Result<float> ExecutiveFitPairs(
     PyMOLGlobals* G, PyObject* list, int quiet);
 
+/**
+ * @param name name of alignment object
+ * @param active_only only consider active alignments
+ * @param state state of alignment object
+ * @return a list of lists of (object, index) tuples containing the
+ * raw per-atom alignment relationships
+ */
+pymol::Result<PyObject*> ExecutiveGetRawAlignment(PyMOLGlobals* G,
+    pymol::null_safe_zstring_view name, bool active_only, int state);
+
 #endif //_PYMOL_NO_PY
diff --git a/layer4/Cmd.cpp b/layer4/Cmd.cpp
@@ -609,79 +609,18 @@ static PyObject *CmdSetRawAlignment(PyObject * self, PyObject * args)
   return APIResult(G, result);
 }
 
-static PyObject* GetRawAlignment(PyMOLGlobals* G,
-    const ObjectAlignment* alnobj,
-    bool active_only,
-    int state)
-{
-  if (state >= alnobj->getNFrame()) {
-    PyErr_Format(PyExc_IndexError, "state %d >= NState %d", state, alnobj->getNFrame());
-    return nullptr;
-  }
-
-  const auto& vla = alnobj->State[state].alignVLA;
-
-  if (!vla) {
-    PyErr_Format(PyExc_IndexError, "state %d not valid", state);
-    return nullptr;
-  }
-
-  auto hide_underscore = SettingGet<bool>(G, cSetting_hide_underscore_names);
-  const auto vla_len = VLAGetSize(vla);
-
-  PyObject * raw = PyList_New(0);
-
-  for (size_t i = 0; i < vla_len; ++i) {
-    PyObject * col = PyList_New(0);
-
-    for (int id; (id = vla[i]); ++i) {
-      auto eoo = ExecutiveUniqueIDAtomDictGet(G, id);
-      if (eoo
-          && (!active_only || eoo->obj->Enabled)
-          && (!hide_underscore || eoo->obj->Name[0] != '_')) {
-        PyObject * idx = Py_BuildValue("si", eoo->obj->Name, eoo->atm + 1);
-        PyList_Append(col, idx);
-        Py_DECREF(idx);
-      }
-    }
-
-    if (PyList_Size(col) > 0) {
-      PyList_Append(raw, col);
-    }
-
-    Py_DECREF(col);
-  }
-
-  return raw;
-}
-
 static PyObject *CmdGetRawAlignment(PyObject * self, PyObject * args)
 {
   PyMOLGlobals *G = NULL;
   const char *name;
   int active_only;
   int state = 0;
-  PyObject *result = NULL;
   API_SETUP_ARGS(G, self, args, "Osi|i", &self, &name, &active_only, &state);
+  pymol::Result<PyObject*> result;
   APIEnterBlocked(G);
-  {
-    if (!name[0]) {
-      name = ExecutiveGetActiveAlignment(G);
-    }
-    if (name && name[0]) {
-      pymol::CObject *obj = ExecutiveFindObjectByName(G, name);
-      if (obj && obj->type == cObjectAlignment) {
-        result = GetRawAlignment(G, (ObjectAlignment*) obj, active_only, state);
-      } else {
-        PyErr_Format(PyExc_KeyError, "no such alignment: '%s'", name);
-      }
-    }
-    APIExitBlocked(G);
-  }
-  if(!result && !PyErr_Occurred()) {
-    return APIFailure(G);
-  } else
-    return result;
+  result = ExecutiveGetRawAlignment(G, name, active_only, state);
+  APIExitBlocked(G);
+  return APIResult(G, result);
 }
 
 static PyObject *CmdGetOrigin(PyObject * self, PyObject * args)