--- a/python/apps/diagrameditor.py	Tue Nov 13 08:16:59 2012 +0100
+++ b/python/apps/diagrameditor.py	Wed Nov 14 18:09:32 2012 +0100
@@ -159,17 +159,15 @@
       return super(PortItem, self).itemChange(change, value)
 
 class OutputPort(PortItem):
-   def __init__(self, name, block):
+   def __init__(self, name, block, d=10.0):
       super(OutputPort, self).__init__(name, block)
-      d = 10.0
       self.setPath(buildPath([QPointF(0.0, -d), QPointF(d, 0), QPointF(0.0, d)]))
    def mousePressEvent(self, event):
       self.scene().startConnection(self)
 
 class InputPort(PortItem):
-   def __init__(self, name, block):
+   def __init__(self, name, block, d=10.0):
       super(InputPort, self).__init__(name, block)
-      d = 10.0
       self.setPath(buildPath([QPointF(-d, -d), QPointF(0, 0), QPointF(-d, d)]))
 
 class HandleItem(QGraphicsEllipseItem):
@@ -190,8 +188,7 @@
       if change == self.ItemPositionChange:
          for cb in self.posChangeCallbacks:
             res = cb(value)
-            if res:
-               value = res
+            if res: value = res
          return value
       return super(HandleItem, self).itemChange(change, value)
 
@@ -235,7 +232,7 @@
       scene = self.scene()
       if scene:
          for view in scene.views():
-            view.dScene = self.subModel
+            view.diagram = self.subModel
             view.zoomAll()
    def newInputPort(self):
       names = [i.name for i in self.inputs + self.outputs]
@@ -253,6 +250,7 @@
       d.update({'name': self.name, 'code': self.code})
       d['inputs'] = [inp.Dict for inp in self.inputs]
       d['outputs'] = [outp.Dict for outp in self.outputs]
+      d['submodel'] = self.subModel.Dict
       return d
    def setDict(self, d):
       self.name = d['name']
@@ -261,6 +259,7 @@
       self.sizer.setPos(d['width'], d['height'])
       for inp in d['inputs']: self.addInput(InputPort(inp['name'], self))
       for outp in d['outputs']: self.addOutput(OutputPort(outp['name'], self))
+      self.subModel.Dict = d['submodel']
    Dict = property(getDict, setDict)
    def addInput(self, i):
       self.inputs.append(i)
@@ -282,8 +281,7 @@
          scene = self.scene()
          if scene:
             for p in self.inputs + self.outputs:
-               if p.connection:
-                  p.connection.myDelete()
+               if p.connection: p.connection.myDelete()
             scene.removeItem(self)
    def updateSize(self):
       rect = self.rect()
@@ -308,30 +306,32 @@
       QGraphicsView.__init__(self, parent)
       self.setDragMode(QGraphicsView.RubberBandDrag)
       self.delShort = QShortcut(QKeySequence.Delete, self)
+      self._model = None
+      self.treeView = QTreeView()
+   def setDiagram(self, d):
+      self.setScene(d)
+      self.delShort.activated.connect(d.deleteItems)
+   def getModel(self): return self._model
    def setModel(self, m):
-      self.setScene(m)
-      self.delShort.activated.connect(m.deleteItems)
-   model = property(lambda s: s.scene(), setModel)
+      self._model = m
+      self.treeView.setModel(m)
+      self.diagram = m.rootDiagram
+   model = property(getModel, setModel)
+   diagram = property(lambda s: s.scene(), setDiagram)
    def save(self):
       filename = 'diagram4.usd'
       with open(filename, 'w') as f:
-         d = self.model.Dict
-         print(d)
-         f.write(json.dumps(d))
+         f.write(json.dumps(self.model.Dict))
    def load(self):
       filename = QFileDialog.getOpenFileName(self)
       self.model = loadModel(filename)
-   def goDown(self):
-      print('Down!')
-      block = self.dScene.selected
    def goUp(self):
-      print('Up!')
-      if hasattr(self.dScene, 'containingBlock'):
-         self.dScene = self.dScene.containingBlock.scene()
+      if hasattr(self.diagram, 'containingBlock'):
+         self.diagram = self.diagram.containingBlock.scene()
          self.zoomAll()
    def zoomAll(self):
       """ zoom to fit all items """
-      rect = self.model.itemsBoundingRect()
+      rect = self.diagram.itemsBoundingRect()
       self.fitInView(rect, Qt.KeepAspectRatio)
    def wheelEvent(self, event):
       pos = event.pos()
@@ -349,33 +349,61 @@
          name = bytes(event.mimeData().data('component/name')).decode()
          pos = self.mapToScene(event.pos())
          self.scene().addNewBlock(pos, name)
+         # TODO: do this in a cleaner way:
+         self.model.modelReset.emit()
 
 class LibraryModel(QStandardItemModel):
-   def __init__(self, parent=None):
-      QStandardItemModel.__init__(self, parent)
    mimeTypes = lambda self: ['component/name']
    def mimeData(self, idxs):
       mimedata = QMimeData()
       for idx in idxs:
          if idx.isValid():
-            txt = self.data(idx, Qt.DisplayRole) # python 3
+            txt = self.data(idx, Qt.DisplayRole)
             mimedata.setData('component/name', txt)
       return mimedata
 
-class ModelHierarchyModel(QStandardItemModel):
+class ModelHierarchyModel(QAbstractItemModel):
    def __init__(self):
       super(ModelHierarchyModel, self).__init__()
       self.rootDiagram = DiagramScene()
-   def flags(self, idx):
-      if idx.isValid():
-         return Qt.ItemIsSelectable | Qt.ItemIsEnabled
-   def data(self):
-      pass
+   def setDict(self, d):
+      self.rootDiagram.Dict = d
+      self.modelReset.emit()
+   Dict = property(lambda s: s.rootDiagram.Dict, setDict)
+   def index(self, row, column, parent=None):
+      if parent.isValid():
+         parent = parent.internalPointer().subModel
+      else:
+         parent = self.rootDiagram
+      blocks = sorted(parent.blocks, key=lambda b:  b.name)
+      block = blocks[row]
+      # Store the index to retrieve it later in the parent function.
+      # TODO: solve this in a better way.
+      block.index = self.createIndex(row, column, block)
+      return block.index
+   def parent(self, index):
+      if index.isValid():
+         block = index.internalPointer()
+         if block.scene() == self.rootDiagram:
+            return QModelIndex()
+         else:
+            outerBlock = block.scene().containingBlock
+            return outerBlock.index
+      print('parent: No valid index')
+   def data(self, index, role):
+      if index.isValid() and role == Qt.DisplayRole:
+         b = index.internalPointer()
+         return b.name
    def headerData(self, section, orientation, role):
       if orientation == Qt.Horizontal and role == Qt.DisplayRole:
-         return "Henkie"
+         return "Element"
    def rowCount(self, parent):
-      pass
+      if parent.column() > 0: return 0
+      if parent.isValid():
+         parent = parent.internalPointer().subModel
+      else:
+         parent = self.rootDiagram
+      return len(parent.blocks)
    def columnCount(self, parent):
       return 1
 
@@ -393,8 +421,7 @@
       for con in d['connections']:
          fromPort = self.findPort(con['fromBlock'], con['fromPort'])
          toPort = self.findPort(con['toBlock'], con['toPort'])
-         c = Connection(fromPort, toPort)
-         self.addItem(c)
+         self.addItem(Connection(fromPort, toPort))
    def getDict(self):
       return {'blocks': [b.Dict for b in self.blocks], 'connections': [c.Dict for c in self.connections]}
    Dict = property(getDict, setDict)
@@ -432,8 +459,7 @@
       self.startedConnection.setEndPos(pos)
       self.addItem(self.startedConnection)
    def deleteItems(self):
-      for item in list(self.selectedItems()):
-         item.myDelete()
+      for item in list(self.selectedItems()): item.myDelete()
 
 class LibraryWidget(QListView):
    def __init__(self):
@@ -455,20 +481,13 @@
       self.setViewMode(self.IconMode)
       self.setDragDropMode(self.DragOnly)
 
-class ModelTree(QTreeView):
-   def __init__(self):
-      super(ModelTree, self).__init__()
-      m = ModelHierarchyModel()
-      self.setModel(m)
-
 def loadModel(filename):
-   with open(filename, 'r') as f:
-      data = f.read()
-   d = json.loads(data)
-   print('Loaded json:', d)
-   s = DiagramScene()
-   s.Dict = d
-   return s
+   m = ModelHierarchyModel()
+   try:
+      with open(filename, 'r') as f: data = f.read()
+      m.Dict = json.loads(data)
+   except FileNotFoundError: pass
+   return m
 
 class Main(QMainWindow):
    def __init__(self):
@@ -486,14 +505,13 @@
       act('Load', QKeySequence.Open, self.editor.load)
       act('Full screen', QKeySequence("F11"), self.toggleFullScreen)
       act('Fit in view', QKeySequence("F8"), self.editor.zoomAll)
-      act('Go down', QKeySequence(Qt.Key_Down), self.editor.goDown)
       act('Go up', QKeySequence(Qt.Key_Up), self.editor.goUp)
       def addDock(name, widget):
          dock = QDockWidget(name, self)
          dock.setWidget(widget)
          self.addDockWidget(Qt.LeftDockWidgetArea, dock)
       addDock('Library', LibraryWidget())
-      addDock('Model tree', ModelTree())
+      addDock('Model tree', self.editor.treeView)
    def toggleFullScreen(self):
       self.setWindowState(self.windowState() ^ Qt.WindowFullScreen)
       self.editor.zoomAll()
@@ -507,6 +525,5 @@
    main.show()
    main.resize(700, 500)
    main.editor.model = loadModel('diagram4.usd')
-   #main.editor.model = DiagramScene()
    app.exec_()