fixed depth_key check in plotting

src/pycea/pl/_utils.py

@@ -13,7 +13,7 @@
 import numpy as np
 from scanpy.plotting import palettes
 
-from pycea.utils import get_leaves, get_root
+from pycea.utils import check_tree_has_key, get_leaves, get_root
 
 
 def layout_nodes_and_branches(
@@ -117,13 +117,12 @@ def layout_trees(
     leaves = []
     depths = []
     for _, tree in trees.items():
+        check_tree_has_key(tree, depth_key)
         tree_leaves = get_leaves(tree)
         leaves.extend(tree_leaves)
         depths.extend(tree.nodes[leaf].get(depth_key) for leaf in tree_leaves)
         if len(depths) != len(leaves):
-            raise ValueError(
-                f"Tree does not have {depth_key} attribute. You can run `pycea.pp.add_depth` to add depth attribute."
-            )
+            raise ValueError(f"Every node in the tree must have a {depth_key} attribute. ")
     max_depth = max(depths)
     n_leaves = len(leaves)
     leaf_coords = {}

src/pycea/utils.py

@@ -37,7 +37,7 @@ def check_tree_has_key(tree: nx.DiGraph, key: str):
     sampled_nodes = random.sample(list(tree.nodes), min(10, len(tree.nodes)))
     for node in sampled_nodes:
         if key not in tree.nodes[node]:
-            message = f"Tree does not have {key} attribute."
+            message = f"Tree nodes to not have {key} attribute."
             if key == "depth":
                 message += " You can run `pycea.pp.add_depth` to add depth attribute."
             raise ValueError(message)