Added ability to optionally return node depth

src/suffix_tree/node.py

@@ -1,6 +1,6 @@
 """Node classes for a Generalized Suffix Tree."""
 
-from typing import Optional, Tuple, List, Set, Dict
+from typing import Optional, Tuple, List, Set, Dict, Union
 
 from .util import Path, Id, Symbol, Symbols, debug
 from . import lca_mixin, util
@@ -204,8 +204,8 @@ def post_order(self, f) -> None:
         f(self)
 
     def find_path(
-        self, S: Symbols, start: int, end: int
-    ) -> Tuple["Node", int, Optional["Node"]]:
+        self, S: Symbols, start: int, end: int, return_nodedepth: bool = False
+    ) -> Union[Tuple["Node", int, Optional["Node"]], Tuple["Node", int, Optional["Node"], int]]:
         """Find a path starting from this node.
 
         The path is absolute.
@@ -217,6 +217,7 @@ def find_path(
         node: Node = self
         matched_len = self.depth
         max_len = end - start
+        node_depth = 0
 
         while matched_len < max_len:
             # find the edge to follow
@@ -231,13 +232,20 @@ def find_path(
                     matched_len += 1
                 if matched_len < child.depth:
                     # the path ends between node and child
+                    if return_nodedepth:
+                        return node, matched_len, None, node_depth
                     return node, matched_len, child
                 # we reached another node, loop
                 node = child
+                node_depth += 1
             else:
                 # no edge to follow
+                if return_nodedepth:
+                    return node, matched_len, None, node_depth
                 return node, matched_len, None
         # path exhausted
+        if return_nodedepth:
+            return node, matched_len, None, node_depth
         return node, matched_len, None
 
     def split_edge(self, new_len: int, child: Node) -> "Internal":

src/suffix_tree/tree.py

@@ -116,15 +116,23 @@ def add(
         assert isinstance(builder, Builder)
         builder.build(self.root, id_, itertools.chain(S, [UniqueEndChar(id_)]))
 
-    def find_path(self, path: Path) -> Tuple[Node, int, Optional[Node]]:
+    def find_path(
+            self, path: Path, return_nodedepth: bool = False,
+            ) -> Union[Tuple[Node, int, Optional[Node]], Tuple[Node, int, Optional[Node], int]]:
         """Find a path in the tree.
 
         See: :py:func:`suffix_tree.node.Internal.find_path`
 
         """
+        if return_nodedepth:
+            return self.root.find_path(
+                path.S, path.start, path.end, return_nodedepth=True
+            )
         return self.root.find_path(path.S, path.start, path.end)
 
-    def find(self, S: Symbols) -> bool:
+    def find(
+            self, S: Symbols, return_nodedepth: bool = False
+        ) -> Union[bool, Tuple[bool, int]]:
         """Find a sequence in the tree.
 
         :param Sequence S: a sequence of symbols
@@ -136,9 +144,18 @@ def find(self, S: Symbols) -> bool:
         True
         >>> tree.find("abc")
         False
+        >>> tree.find("abx", return_depth=True)
+        (True, 2)
         """
 
         path = Path(S)
+
+        if return_nodedepth:
+            dummy_node, matched_len, dummy_child, nodedepth = self.find_path(
+                path, return_nodedepth=return_nodedepth
+            )
+            return matched_len == len(path), nodedepth
+
         dummy_node, matched_len, dummy_child = self.find_path(path)
         return matched_len == len(path)