Make item additions undo-able

1 files changed, 24 insertions, 15 deletions

diff --git a/src/diagramdocument.cpp b/src/diagramdocument.cpp
index fe8c53f..d0a8b4c 100644
--- a/src/diagramdocument.cpp
+++ b/src/diagramdocument.cpp
@@ -52,6 +52,18 @@ DiagramDocument::updatePositions(DatabaseTable *table)
 	}
 }
 
+template <class T> QList<T *>
+DiagramDocument::itemsByType()
+{
+	QList<T *> result;
+	foreach(QGraphicsItem *item, items()) {
+		if (item->type() == T::Type) {
+			result.append(static_cast<T *>(item));
+		}
+	}
+	return result;
+}
+
 void
 DiagramDocument::itemMoved(DiagramItem *item)
 {
@@ -77,8 +89,7 @@ DiagramDocument::mousePressEvent(QGraphicsSceneMouseEvent *event)
 		DatabaseTable *table = new DatabaseTable();
 		table->setPos(event->scenePos());
 		table->setZValue(10.0);
-		m_tables << table;
-		addItem(table);
+		undoStack()->push(new AddItemCommand(this, table));
 		clearSelection();
 		table->setSelected(true);
 		setMode(Select);
@@ -121,7 +132,6 @@ DiagramDocument::mouseReleaseEvent(QGraphicsSceneMouseEvent *event)
 			relation->setTarget(target);
 			relation->setZValue(1.0);
 			addItem(relation);
-			m_relations << relation;
 		}
 
 		delete m_line;
@@ -157,10 +167,8 @@ DiagramDocument::deleteSelectedItems()
 		if (table) {
 			foreach (DatabaseRelationship *relation, findTableRelations(table)) {
 				removeItem(relation);
-				m_relations.removeAll(relation);
 			}
 			removeItem(table);
-			m_tables.removeAll(table);
 		}
 	}
 }
@@ -169,7 +177,7 @@ QList<DatabaseRelationship *>
 DiagramDocument::findTableRelations(DatabaseTable *table)
 {
 	QList<DatabaseRelationship *> result;
-	foreach (DatabaseRelationship *relation, m_relations) {
+	foreach (DatabaseRelationship *relation, itemsByType<DatabaseRelationship>()) {
 		if (relation->source() == table || relation->target() == table) {
 			result << relation;
 		}
@@ -208,7 +216,8 @@ DiagramDocument::save(const QString &fileName)
 	QDomElement tableListElement = doc.createElement("table-list");
 	root.appendChild(tableListElement);
 
-	foreach (DatabaseTable *table, m_tables) {
+	QList<DatabaseTable *> tables = itemsByType<DatabaseTable>();
+	foreach (DatabaseTable *table, tables) {
 		QDomElement tableElement = doc.createElement("table");
 		tableListElement.appendChild(tableElement);
 		
@@ -238,10 +247,10 @@ DiagramDocument::save(const QString &fileName)
 	QDomElement relationListElement = doc.createElement("relation-list");
 	root.appendChild(relationListElement);
 
-	foreach (DatabaseRelationship *relation, m_relations) {
+	foreach (DatabaseRelationship *relation, itemsByType<DatabaseRelationship>()) {
 		QDomElement relationElement = doc.createElement("relation");
-		relationElement.setAttribute("from", QString::number(m_tables.indexOf(relation->sourceTable())));
-		relationElement.setAttribute("to", QString::number(m_tables.indexOf(relation->targetTable())));
+		relationElement.setAttribute("from", QString::number(tables.indexOf(relation->sourceTable())));
+		relationElement.setAttribute("to", QString::number(tables.indexOf(relation->targetTable())));
 		relationListElement.appendChild(relationElement);
 	}
 
@@ -269,6 +278,8 @@ DiagramDocument::load(const QString &fileName)
 
 	QDomElement root = doc.firstChildElement("database-model");
 
+	QList<DatabaseTable *> tables;
+
 	QDomElement tableListElement = root.firstChildElement("table-list");
 	QDomElement tableElement = tableListElement.firstChildElement("table");
 	while (!tableElement.isNull()) {
@@ -297,8 +308,8 @@ DiagramDocument::load(const QString &fileName)
 			columnElement = columnElement.nextSiblingElement("column");
 		}
 
-		m_tables << table;
 		addItem(table);
+		tables.append(table);
 
 		tableElement = tableElement.nextSiblingElement("table");
 	}
@@ -313,12 +324,10 @@ DiagramDocument::load(const QString &fileName)
 			int index1 = relationElement.attribute("to").toInt(&ok);
 			if (ok) {
 				DatabaseRelationship *relation = new DatabaseRelationship();
-				relation->setSource(m_tables[index0]);
-				relation->setTarget(m_tables[index1]);
-				m_relations << relation;
+				relation->setSource(tables[index0]);
+				relation->setTarget(tables[index1]);
 				addItem(relation);
 			}
 		}
 	}
-
 }