Implement node restoring

plasma/layout.py

@@ -3,7 +3,7 @@
 from xcffib.xproto import StackMode
 from libqtile.layout.base import Layout
 
-from .node import Node, AddMode
+from .node import Node, AddMode, NotRestorableError
 
 
 class Plasma(Layout):
@@ -58,12 +58,12 @@ def clone(self, group):
         return clone
 
     def add(self, client):
-        current = self.focused_node
-        new_node = Node(client)
-        if self.focused is None or current is None:
-            self.root.add_child(new_node)
-        else:
-            current.add_node(new_node, self.add_mode)
+        node = self.focused_node or self.root
+        new = Node(client)
+        try:
+            self.root.restore(new)
+        except NotRestorableError:
+            node.add_node(new, self.add_mode)
         self.add_mode = None
 
     def remove(self, client):

plasma/node.py

@@ -45,6 +45,9 @@ def orient(self):
     RIGHT: lambda a, b: isclose(a.x_end, b.x),
 }
 
+class NotRestorableError(Exception):
+    pass
+
 class Node:
     """A tree node.
 
@@ -63,13 +66,22 @@ def __init__(self, payload=None, x=None, y=None, width=None, height=None):
         self.children = []
         self.last_accessed = datetime.min
         self.parent = None
+        self.restorables = {}
 
     def __repr__(self):
         info = self.payload or ''
         if self.children:
             info += ' +%d' % len(self.children)
         return '<Node %s %x>' % (info, id(self))
 
+    def __contains__(self, node):
+        if node is self:
+            return True
+        for child in self.children:
+            if node in child:
+                return True
+        return False
+
     @property
     def root(self):
         try:
@@ -292,18 +304,22 @@ def size(self):
     def size(self, val):
         if self.is_root or not self.siblings:
             return
+        if val is None:
+            self.reset_size()
+            return
         occupied = sum(s.min_size_bound for s in self.siblings)
         val = max(min(val, self.parent.capacity - occupied),
                   self.min_size_bound)
         self.force_size(val)
 
     def force_size(self, val):
         """Set size without considering available space."""
-        self.fit_into(self.siblings, self.parent.capacity - val)
-        if val != 0:
-            if self.children:
-                self.fit_into([self], val)
-            self._size = val
+        Node.fit_into(self.siblings, self.parent.capacity - val)
+        if val == 0:
+            return
+        if self.children:
+            Node.fit_into([self], val)
+        self._size = val
 
     @property
     def size_offset(self):
@@ -365,7 +381,7 @@ def reset_size(self):
         self._size = None
 
     def grow(self, amt, orient=None):
-        # TODO Deprecate grow
+        # TODO Deprecate grow, it should be replaced by size assignment
         if self.is_root:
             return
         if orient is ~self.parent.orient:
@@ -463,29 +479,30 @@ def close_right(self):
         return self.close_neighbor(RIGHT)
 
     def add_child(self, node, idx=None):
+        # TODO Currently we assume that the node has no size.
         if idx is None:
             idx = len(self.children)
         self.children.insert(idx, node)
         node.parent = self
         if len(self.children) == 1:
             return
         total = self.capacity
-        self.fit_into(node.siblings, total - (total / len(self.children)))
+        Node.fit_into(node.siblings, total - (total / len(self.children)))
 
     def add_child_after(self, new, old):
         self.add_child(new, idx=self.children.index(old)+1)
 
     def remove_child(self, node):
+        node._save_restore_state()  # pylint: disable=W0212
         node.force_size(0)
         self.children.remove(node)
-        if len(self.children) != 1:
-            return
-        if not self.is_root:
-            # Collapse tree with a single child
-            self.parent.replace_child(self, self.children[0])
-        else:
-            # A single child doesn't need an absolute size
-            self.children[0].reset_size()
+        if len(self.children) == 1:
+            if self.is_root:
+                # A single child doesn't need a fixed size
+                self.children[0].reset_size()
+            else:
+                # Collapse tree with a single child
+                self.parent.replace_child(self, self.children[0])
 
     def remove(self):
         self.parent.remove_child(self)
@@ -495,19 +512,19 @@ def replace_child(self, old, new):
         new.parent = self
         new._size = old._size  # pylint: disable=protected-access
 
-    def flip_with(self, node):
+    def flip_with(self, node, reverse=False):
         """Join with node in a new, orthogonal container."""
         container = Node()
         self.parent.replace_child(self, container)
         self.reset_size()
-        for child in [self, node]:
+        for child in [node, self] if reverse else [self, node]:
             container.add_child(child)
 
     def add_node(self, node, mode=None):
         """Add node according to the mode.
 
-        This can result in adding it as a child, joining with it in a new,
-        flipped sub-container or splitting the space with it.
+        This can result in adding it as a child, joining with it in a new
+        flipped sub-container, or splitting the space with it.
         """
         if self.is_root:
             self.add_child(node)
@@ -523,6 +540,41 @@ def add_node(self, node, mode=None):
         else:
             self.flip_with(node)
 
+    def restore(self, node):
+        """Restore node.
+
+        Try to add the node in a place where a node with the same payload
+        has previously been.
+        """
+        try:
+            parent, idx, sizes, flip = self.root.restorables[node.payload]
+        except KeyError:
+            raise NotRestorableError()
+        if parent not in self.root:
+            # Don't restore at a parent that's not part of the tree anymore
+            raise NotRestorableError()
+        node.reset_size()
+        if flip:
+            parent.flip_with(node, reverse=(idx == 0))
+            node.size, parent.size = sizes
+        else:
+            parent.add_child(node, idx=idx)
+            node.size = sizes[0]
+            if len(sizes) == 2:
+                node.siblings[0].size = sizes[1]
+        del self.root.restorables[node.payload]
+
+    def _save_restore_state(self):
+        parent = self.parent
+        sizes = (self._size,)
+        flip = False
+        if len(self.siblings) == 1:
+            sizes += (self.siblings[0]._size,)  # pylint: disable=W0212
+            if not self.parent.is_root:
+                flip = True
+                parent = self.siblings[0]
+        self.root.restorables[self.payload] = (parent, self.index, sizes, flip)
+
     def move(self, direction):
         if self.is_root:
             return

tests/test_plasma.py

@@ -1,7 +1,8 @@
+import pytest
 from pytest import approx
 
 from plasma.debug import draw, tree, info
-from plasma.node import Node, VERTICAL, HORIZONTAL, AddMode
+from plasma.node import Node, VERTICAL, HORIZONTAL, AddMode, NotRestorableError
 
 from .conftest import Nodes
 
@@ -866,12 +867,10 @@ def test_add_node(self, root):
         a.add_node(c)
         assert root.tree == [a, c, b]
         c.add_node(d, mode=AddMode.HORIZONTAL)
-        info(root)
         assert root.tree == [a, c, d, b]
         root.remove_child(d)
         c.add_node(d, mode=AddMode.VERTICAL)
         c.parent.add_child_after
-        info(root)
         assert root.tree == [a, [c, d], b]
         c.add_node(e, mode=AddMode.VERTICAL)
         assert root.tree == [a, [c, e, d], b]
@@ -883,6 +882,131 @@ def test_add_node(self, root):
         a.add_node(g, mode=AddMode.VERTICAL | AddMode.SPLIT)
         assert root.tree == [[a, g], f, [c, e, d], b]
 
+    def test_contains(self, root, grid):
+        x = Node('x')
+        nodes = list(grid)
+        nodes += [n.parent for n in nodes]
+        nodes.append(root)
+        for n in nodes:
+            assert n in root
+        assert x not in root
+
+    def test_restore(self, root, grid):
+        a, b, c, d, e = grid
+        tree = root.tree
+        for node in grid:
+            node.remove()
+            root.restore(node)
+            assert root.tree == tree
+
+    def test_restore_same_payload(self, root, grid):
+        """Restore a node that's not identical with the removed one but carries
+        the same payload.
+        """
+        a, b, c, d, e = grid
+        d.remove()
+        new_d = Node('d')
+        root.restore(new_d)
+        assert root.tree == [a, [b, [c, new_d, e]]]
+
+    def test_restore_unknown(self, root, grid):
+        a, b, c, d, e = grid
+        with pytest.raises(NotRestorableError):
+            root.restore(Node('x'))
+        d.remove()
+        with pytest.raises(NotRestorableError):
+            root.restore(Node('x'))
+        root.restore(d)
+        assert root.tree == [a, [b, [c, d, e]]]
+
+    def test_restore_no_parent(self, root, small_grid):
+        a, b, c, d = small_grid
+        c.remove()
+        d.remove()
+        with pytest.raises(NotRestorableError):
+            root.restore(c)
+        root.restore(d)
+        assert root.tree == [a, [b, d]]
+
+    def test_restore_bad_index(self, root, grid):
+        a, b, c, d, e = grid
+        f, g = Nodes('f g')
+        e.parent.add_child(f)
+        e.parent.add_child(g)
+        g.remove()
+        f.remove()
+        e.remove()
+        root.restore(g)
+        assert root.tree == [a, [b, [c, d, g]]]
+
+    def test_restore_sizes(self, root, grid):
+        a, b, c, d, e = grid
+        c.size = 30
+        c.remove()
+        root.restore(c)
+        assert c.size == 30
+        c.remove()
+        d.size = 30
+        e.size = 30
+        assert d.size == e.size == 30
+        root.restore(c)
+        assert c.size == 30
+        assert d.size == e.size == 15
+
+    def test_restore_sizes_flip(self, root, tiny_grid):
+        a, b, c = tiny_grid
+        c.size = 10
+        c.remove()
+        assert a._size is b._size is None
+        root.restore(c)
+        assert c.size == 10
+        b.size = 10
+        c.remove()
+        root.restore(c)
+        assert b.size == 10
+        b.remove()
+        root.restore(b)
+        assert b.size == 10
+
+    def test_restore_root(self, root):
+        a, b = Nodes('a b')
+        root.add_child(a)
+        root.add_child(b)
+        a.size = 20
+        a.remove()
+        root.restore(a)
+        assert a._size == 20
+        assert b._size is None
+        b.remove()
+        root.restore(b)
+        assert a._size == 20
+        assert b._size is None
+
+    def test_restore_keep_flexible(self, root, tiny_grid):
+        a, b, c = tiny_grid
+        b.remove()
+        root.restore(b)
+        assert a._size is b._size is c._size is None
+        b.size = 10
+        b.remove()
+        root.restore(b)
+        assert b._size == 10
+        assert c._size is None
+        c.remove()
+        root.restore(c)
+        assert b._size == 10
+        assert c._size is None
+        c.size = 10
+        b.reset_size()
+        b.remove()
+        root.restore(b)
+        assert b._size is None
+        assert c._size == 10
+        c.remove()
+        root.restore(c)
+        assert b._size is None
+        assert c._size == 10
+
 class TestDebugging:
 
     def test_tree(self, root, grid):