Repository: scalableminds/webknossos-libs

Issue:
Adding a tree from one skeleton to another does not work
```python
import webknossos as wk

import webknossos as wk

def main() -> None:
    skeleton_a = wk.Skeleton.load("./a.nml")
    skeleton_b = wk.Skeleton.load("./b.nml")

    for tree in skeleton_a.flattened_trees():
        skeleton_b.add_tree(tree)

    skeleton_b.save("./c.nml")

if __name__ == '__main__':
    main()

```

produces only an nml with empty trees (the `<thing>` elements are there, but nodes and edges are missing).


Discussion:
@philippotto The code example above will only add the tree to `skeleton_b` in the error case, i.e. when it was not found. So yes, the resulting new NML should be empty.
The original logic of the code is to copy a tree from skeleton A to skeleton B if it's not present in B. So, in the end, B should contain all trees from A. The code even works with regard to copying the high level tree object, but the trees themselves are empty then (i.e., the nodes and edges are missing).

I simplified the code a bit, since the try-except wasn't really necessary to exemplify the bug.

Resolving patch:
diff --git a/webknossos/examples/merge_trees_of_nml_files.py b/webknossos/examples/merge_trees_of_nml_files.py
new file mode 100644
--- /dev/null
+++ b/webknossos/examples/merge_trees_of_nml_files.py
@@ -0,0 +1,15 @@
+import webknossos as wk
+
+
+def main() -> None:
+    skeleton_a = wk.Skeleton.load("./a.nml")
+    skeleton_b = wk.Skeleton.load("./b.nml")
+
+    for tree in skeleton_a.flattened_trees():
+        skeleton_b.add_tree(tree)
+
+    skeleton_b.save("./c.nml")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/webknossos/webknossos/skeleton/group.py b/webknossos/webknossos/skeleton/group.py
--- a/webknossos/webknossos/skeleton/group.py
+++ b/webknossos/webknossos/skeleton/group.py
@@ -1,3 +1,4 @@
+import copy
 from typing import TYPE_CHECKING, Iterator, Optional, Set, Tuple, Union, cast
 
 import attr
@@ -63,25 +64,53 @@ def id(self) -> int:
 
     def add_tree(
         self,
-        name: str,
+        name_or_tree: Union[str, Tree],
         color: Optional[Union[Vector4, Vector3]] = None,
         _enforced_id: Optional[int] = None,
     ) -> Tree:
-        """Adds a tree to the current group with the provided name (and color if specified)."""
+        """Adds a tree to the current group. If the first parameter is a string,
+        a new tree will be added with the provided name and color if specified.
+        Otherwise, the first parameter is assumed to be a tree object (e.g., from
+        another skeleton). A copy of that tree will then be added. If the id
+        of the tree already exists, a new id will be generated."""
 
         if color is not None and len(color) == 3:
             color = cast(Optional[Vector4], color + (1.0,))
         color = cast(Optional[Vector4], color)
-        new_tree = Tree(
-            name=name,
-            color=color,
-            group=self,
-            skeleton=self._skeleton,
-            enforced_id=_enforced_id,
-        )
-        self._child_trees.add(new_tree)
 
-        return new_tree
+        if type(name_or_tree) is str:
+            name = name_or_tree
+            new_tree = Tree(
+                name=name,
+                color=color,
+                group=self,
+                skeleton=self._skeleton,
+                enforced_id=_enforced_id,
+            )
+            self._child_trees.add(new_tree)
+
+            return new_tree
+        else:
+            tree = cast(Tree, name_or_tree)
+            new_tree = copy.deepcopy(tree)
+
+            if color is not None:
+                new_tree.color = color
+
+            if _enforced_id is not None:
+                assert not self.has_tree_id(
+                    _enforced_id
+                ), "A tree with the specified _enforced_id already exists in this group."
+                new_tree._id = _enforced_id
+
+            if self.has_tree_id(tree.id):
+                new_tree._id = self._skeleton._element_id_generator.__next__()
+
+            new_tree.group = self
+            new_tree.skeleton = self._skeleton
+
+            self._child_trees.add(new_tree)
+            return new_tree
 
     def add_graph(
         self,
@@ -91,7 +120,10 @@ def add_graph(
     ) -> Tree:
         """Deprecated, please use `add_tree`."""
         warn_deprecated("add_graph()", "add_tree()")
-        return self.add_tree(name=name, color=color, _enforced_id=_enforced_id)
+        return self.add_tree(name_or_tree=name, color=color, _enforced_id=_enforced_id)
+
+    def remove_tree_by_id(self, tree_id: int) -> None:
+        self._child_trees.remove(self.get_tree_by_id(tree_id))
 
     @property
     def children(self) -> Iterator[GroupOrTree]:
@@ -180,6 +212,14 @@ def get_tree_by_id(self, tree_id: int) -> Tree:
                 return tree
         raise ValueError(f"No tree with id {tree_id} was found")
 
+    def has_tree_id(self, tree_id: int) -> bool:
+        """Returns true if this group (or a subgroup) contains a tree with the given id."""
+        try:
+            self.get_tree_by_id(tree_id)
+            return True
+        except ValueError:
+            return False
+
     def get_graph_by_id(self, graph_id: int) -> Tree:
         """Deprecated, please use `get_tree_by_id`."""
         warn_deprecated("get_graph_by_id()", "get_tree_by_id()")
diff --git a/webknossos/webknossos/skeleton/tree.py b/webknossos/webknossos/skeleton/tree.py
--- a/webknossos/webknossos/skeleton/tree.py
+++ b/webknossos/webknossos/skeleton/tree.py
@@ -114,7 +114,7 @@ def __init__(
         group: "Group",
         skeleton: "Skeleton",
         color: Optional[Vector4] = None,
-        enforced_id: Optional[int] = None,
+        enforced_id: Optional[int] = None,  # pylint: disable=unused-argument
     ) -> None:
         """
         To create a tree, it is recommended to use `Skeleton.add_tree` or
@@ -122,37 +122,64 @@ def __init__(
         attached as a child to the object the method was called on.
         """
 
-        # To be able to reference nodes by id after adding them for the first time, we use custom dict-like classes
-        # for the networkx-graph structures, that have nodes as keys:
-        # * `self._node`: _NodeDict
-        #   holding the attributes of nodes, keeping references from ids to nodes
-        # * `self._adj`: _AdjDict on the first two levels
-        #   holding edge attributes on the last level, using self._node to convert ids to nodes
-        #
-        # It's important to set the attributes before the parent's init so that they shadow the class-attributes.
-        #
-        # For further details, see the *Subclasses* section here: https://networkx.org/documentation/stable/reference/classes/graph.html
-
-        self.node_dict_factory = _NodeDict
-        # The lambda works because self._node is set before self._adj in networkx.Graph.__init__
-        # and because the lambda is evaluated lazily.
-        self.adjlist_outer_dict_factory = lambda: _AdjDict(node_dict=self._node)
-        self.adjlist_inner_dict_factory = lambda: _AdjDict(node_dict=self._node)
-
         super().__init__()
-
+        # Note that id is set up in __new__
         self.name = name
         self.group = group
         self.color = color
 
-        # read-only member, exposed via properties
+        # only used internally
+        self._skeleton = skeleton
+
+    def __new__(
+        cls,
+        name: str,  # pylint: disable=unused-argument
+        group: "Group",  # pylint: disable=unused-argument
+        skeleton: "Skeleton",  # pylint: disable=unused-argument
+        color: Optional[Vector4] = None,  # pylint: disable=unused-argument
+        enforced_id: Optional[int] = None,
+    ) -> "Tree":
+        self = super().__new__(cls)
+
+        # self._id is a read-only member, exposed via properties.
+        # It is set in __new__ instead of __init__ so that pickling/unpickling
+        # works without problems. As long as the deserialization of a tree instance
+        # is not finished, the object is only half-initialized. Since self._id
+        # is needed by __hash__, an error would be raised otherwise.
+        # Also see:
+        # https://stackoverflow.com/questions/46283738/attributeerror-when-using-python-deepcopy
         if enforced_id is not None:
             self._id = enforced_id
         else:
             self._id = skeleton._element_id_generator.__next__()
 
-        # only used internally
-        self._skeleton = skeleton
+        return self
+
+    def __getnewargs__(self) -> Tuple:
+        # pickle.dump will pickle instances of Tree so that the following
+        # tuple is passed as arguments to __new__.
+        return (self.name, self.group, self._skeleton, self.color, self._id)
+
+    # node_dict_factory, adjlist_outer_dict_factory and adjlist_inner_dict_factory are used by networkx
+    # from which we subclass.
+    # To be able to reference nodes by id after adding them for the first time, we use custom dict-like classes
+    # for the networkx-graph structures, that have nodes as keys:
+    #     * `self._node`: _NodeDict
+    #        holding the attributes of nodes, keeping references from ids to nodes
+    #     * `self._adj`: _AdjDict on the first two levels
+    #       holding edge attributes on the last level, using self._node to convert ids to nodes
+    # It's important to set the attributes before the parent's init so that they shadow the class-attributes.
+    # For further details, see the *Subclasses* section here: https://networkx.org/documentation/stable/reference/classes/graph.html
+
+    node_dict_factory = _NodeDict
+
+    def adjlist_outer_dict_factory(self) -> _AdjDict:
+        # self._node will already be available when this method is called, because networkx.Graph.__init__
+        # sets up the nodes first and then the edges (i.e., adjacency list).
+        return _AdjDict(node_dict=self._node)
+
+    def adjlist_inner_dict_factory(self) -> _AdjDict:
+        return _AdjDict(node_dict=self._node)
 
     def __to_tuple_for_comparison(self) -> Tuple:
         return (


Test patch:
diff --git a/webknossos/tests/test_skeleton.py b/webknossos/tests/test_skeleton.py
--- a/webknossos/tests/test_skeleton.py
+++ b/webknossos/tests/test_skeleton.py
@@ -275,3 +275,72 @@ def test_load_nml(tmp_path: Path) -> None:
     skeleton_a = wk.Skeleton.load(input_path)
     skeleton_a.save(output_path)
     assert skeleton_a == wk.Skeleton.load(output_path)
+
+
+def test_remove_tree(tmp_path: Path) -> None:
+    input_path = TESTDATA_DIR / "nmls" / "test_a.nml"
+    output_path = tmp_path / "test_a.nml"
+    skeleton_a = wk.Skeleton.load(input_path)
+
+    # Check that tree exists
+    tree = skeleton_a.get_tree_by_id(1)
+    assert tree is not None
+    assert tree in list(skeleton_a.children)
+
+    # Check that tree doesn't exist anymore
+    skeleton_a.remove_tree_by_id(1)
+    with pytest.raises(ValueError):
+        tree = skeleton_a.get_tree_by_id(1)
+
+    assert tree not in list(skeleton_a.children)
+
+    # Check that serialized skeleton doesn't contain
+    # deleted tree
+    skeleton_a.save(output_path)
+    assert skeleton_a == wk.Skeleton.load(output_path)
+
+    # Load original file and check that tree is still
+    # there (should not have been removed on disk
+    # automatically).
+    skeleton_a = wk.Skeleton.load(input_path)
+    assert tree in list(skeleton_a.children)
+
+
+def test_add_tree_with_obj(tmp_path: Path) -> None:
+    input_path = TESTDATA_DIR / "nmls" / "test_a.nml"
+    output_path = tmp_path / "test_a.nml"
+    skeleton_a = wk.Skeleton.load(input_path)
+
+    # Check that tree exists
+    tree = skeleton_a.get_tree_by_id(1)
+
+    skeleton_b = wk.Skeleton(
+        voxel_size=skeleton_a.voxel_size, dataset_name=skeleton_a.dataset_name
+    )
+    skeleton_b.add_tree(tree)
+
+    assert tree is not skeleton_b.get_tree_by_id(1)
+    assert tree == skeleton_b.get_tree_by_id(1)
+
+    skeleton_b.save(output_path)
+
+
+def test_add_tree_with_obj_and_properties(tmp_path: Path) -> None:
+    input_path = TESTDATA_DIR / "nmls" / "test_a.nml"
+    output_path = tmp_path / "test_a.nml"
+    skeleton_a = wk.Skeleton.load(input_path)
+
+    # Check that tree exists
+    tree = skeleton_a.get_tree_by_id(1)
+
+    skeleton_b = wk.Skeleton(
+        voxel_size=skeleton_a.voxel_size, dataset_name=skeleton_a.dataset_name
+    )
+    new_tree = skeleton_b.add_tree(tree, color=(1, 2, 3), _enforced_id=1337)
+
+    assert new_tree is skeleton_b.get_tree_by_id(1337)
+    assert new_tree is not tree
+    assert new_tree != tree
+    assert new_tree.color == (1, 2, 3, 1)
+
+    skeleton_b.save(output_path)