Add remove_tree and overload add_tree (#891)

* first draft for remove_tree and add_tree with tree object

* use overloading for add_tree

* format

* format

* trigger ci

* fix linting

* format

* fix typing

* clean up

* beautify comment

* also handle color and _enforced_id and test this

* update changelog

* add example script for merging NMLs

* format

* fix linting

* move example file to correct folder

Author: philippotto

docs/src/webknossos-py/examples/merge_nmls.md

@@ -0,0 +1,9 @@
+# Merge Trees of two NML files
+
+This example opens two [NML files](/webknossos/data_formats.html#nml-files) and copies the trees from the NML A to B and saves the output to a new NML (that includes the trees of both input NMLs).
+
+```python
+--8<--
+webknossos/examples/merge_trees_of_nml_files.py
+--8<--
+```

webknossos/Changelog.md

@@ -15,6 +15,8 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 ### Breaking Changes
 
 ### Added
+- `Group.add_tree` now also accepts a tree object as a first parameter (instead of only a string). [#891](https://github.com/scalableminds/webknossos-libs/pull/891)
+- `Group.remove_tree_by_id` was added. [#891](https://github.com/scalableminds/webknossos-libs/pull/891)
 
 ### Changed
 - Upgrades `black`, `mypy`, `pylint`, `pytest`. [#873](https://github.com/scalableminds/webknossos-libs/pull/873)

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()

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)

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()")

webknossos/webknossos/skeleton/tree.py

@@ -114,45 +114,72 @@ 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
         `Group.add_tree`. That way, the newly created tree is automatically
         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 (