Merge pull request #632 from ghutchis/more-selection-commands

More selection commands
OpenChemistry · Jun 26, 2021 · 9351bb0 · 9351bb0
2 parents 1420315 + 86bc032
commit 9351bb0
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 1 deletion.
diff --git a/avogadro/qtplugins/select/select.cpp b/avogadro/qtplugins/select/select.cpp
@@ -17,9 +17,11 @@
 #include "select.h"
 
 #include <avogadro/qtgui/molecule.h>
+#include <avogadro/qtgui/periodictableview.h>
 
 #include <QtGui/QKeySequence>
 #include <QtWidgets/QAction>
+#include <QtWidgets/QInputDialog>
 
 #include <QtCore/QStringList>
 
@@ -29,7 +31,8 @@ namespace Avogadro {
 namespace QtPlugins {
 
 Select::Select(QObject* parent_)
-  : Avogadro::QtGui::ExtensionPlugin(parent_), m_molecule(nullptr)
+  : Avogadro::QtGui::ExtensionPlugin(parent_), m_molecule(nullptr),
+    m_elements(nullptr)
 {
   QAction* action = new QAction(tr("Select All"), this);
   action->setShortcut(QKeySequence("Ctrl+A"));
@@ -48,10 +51,20 @@ Select::Select(QObject* parent_)
   action = new QAction(tr("Invert Selection"), this);
   connect(action, SIGNAL(triggered()), SLOT(invertSelection()));
   m_actions.append(action);
+
+  action = new QAction(tr("Select by Element..."), this);
+  connect(action, SIGNAL(triggered()), SLOT(selectElement()));
+  m_actions.append(action);
+
+  action = new QAction(tr("Select by Atom Index..."), this);
+  connect(action, SIGNAL(triggered()), SLOT(selectAtomIndex()));
+  m_actions.append(action);
 }
 
 Select::~Select()
 {
+  if (m_elements)
+    m_elements->deleteLater();
 }
 
 QString Select::description() const
@@ -94,6 +107,72 @@ void Select::selectNone()
   }
 }
 
+void Select::selectElement()
+{
+  if (!m_molecule)
+    return;
+
+  if (m_elements == nullptr) {
+    m_elements = new QtGui::PeriodicTableView(qobject_cast<QWidget*>(parent()));
+    connect(m_elements, SIGNAL(elementChanged(int)), this,
+            SLOT(selectElement(int)));
+  }
+
+  m_elements->show();
+}
+
+void Select::selectElement(int element)
+{
+  if (!m_molecule)
+    return;
+
+  for (Index i = 0; i < m_molecule->atomCount(); ++i) {
+    if (m_molecule->atomicNumber(i) == element)
+      m_molecule->atom(i).setSelected(true);
+    else
+      m_molecule->atom(i).setSelected(false);
+  }
+
+  m_molecule->emitChanged(Molecule::Atoms);
+}
+
+void Select::selectAtomIndex()
+{
+  if (!m_molecule)
+    return;
+
+  bool ok;
+  QString text = QInputDialog::getText(
+    qobject_cast<QWidget*>(parent()), tr("Select Atoms by Index"),
+    tr("Atoms to Select:"), QLineEdit::Normal, QString(), &ok);
+
+  if (!ok || text.isEmpty())
+    return;
+
+  auto list = text.simplified().split(',');
+  foreach (QString item, list) {
+    // check if it's a range
+    if (item.contains('-')) {
+      auto range = item.split('-');
+      if (range.size() >= 2) {
+        bool ok1, ok2;
+        int start = range.first().toInt(&ok1);
+        int last = range.back().toInt(&ok2);
+        if (ok1 && ok2) {
+          for (Index i = start; i <= last; ++i)
+            m_molecule->atom(i).setSelected(true);
+        }
+      }
+    } else {
+      int i = item.toInt(&ok);
+      if (ok)
+        m_molecule->atom(i).setSelected(true);
+    }
+  }
+
+  m_molecule->emitChanged(Molecule::Atoms);
+}
+
 void Select::invertSelection()
 {
   if (m_molecule) {

diff --git a/avogadro/qtplugins/select/select.h b/avogadro/qtplugins/select/select.h
@@ -20,6 +20,10 @@
 #include <avogadro/qtgui/extensionplugin.h>
 
 namespace Avogadro {
+
+namespace QtGui{
+class PeriodicTableView;
+}
 namespace QtPlugins {
 
 /**
@@ -44,10 +48,14 @@ private slots:
   void selectAll();
   void selectNone();
   void invertSelection();
+  void selectElement();
+  void selectAtomIndex();
+  void selectElement(int element);
 
 private:
   QList<QAction*> m_actions;
   QtGui::Molecule* m_molecule;
+  QtGui::PeriodicTableView* m_elements;
 };
 
 } // namespace QtPlugins