Introduce walk_towards method for walking a path to node with given key

Provide a walk_towards method which generalises prefixes method and
allows walking towards particular node and accessing each step of the
path.

Compared to prefixes method, steps for nodes without assigned values
are returned and the method raises KeyError if trying to walk towards
non-existent node.

Furthermore, at each step value of a node can be manipulated with
set and setdefault methods.

Fixes: google#20

pygtrie.py

@@ -275,13 +275,6 @@ def __setstate__(self, state):
                 stack[-1].value = next(state)
 
 
-_NONE_PAIR = type('NonePair', (tuple,), {
-    '__nonzero__': lambda _: False,
-    '__bool__': lambda _: False,
-    '__slots__': (),
-})((None, None))
-
-
 class Trie(_collections.MutableMapping):
     """A trie implementation with dict interface plus some extensions.
 
@@ -878,6 +871,130 @@ def __delitem__(self, key_or_slice):
         node.value = _SENTINEL
         self._cleanup_trace(trace)
 
+    class _NoneStep(object):
+        """Representation of a non-existent step towards non-existent node."""
+
+        __slots__ = ()
+
+        __bool__ = __nonzero__ = lambda self: False
+        is_set = has_subtrie = property(lambda self: False)
+        get = lambda self, default=None: default
+
+        key = value = property(lambda self: None)
+
+        def __getitem__(self, index):
+            """Makes object appear like a (key, value) tuple.
+
+            This is deprecated and for backwards-compatibility only.  Prefer
+            using ``key`` and ``value`` properties directly.
+
+            Args:
+                index: Element index to return.  Zero for key, one for value.
+
+            Returns:
+                ``self.key`` if index is ``0``, ``self.value`` if it's ``1``.
+                Otherwise raises an IndexError exception.
+
+            Raises:
+                IndexError: if index is not 0 or 1.
+                KeyErro: if index is 1 but node has no value assigned.
+            """
+            if index == 0:
+                return self.key
+            elif index == 1:
+                return self.value
+            else:
+                raise IndexError('index out of range')
+
+    class _Step(_NoneStep):
+        """Representation of a single step on a path towards particular node."""
+
+        __slots__ = ('_trie', '_path', '_pos', '_node', '__key')
+
+        def __init__(self, trie, path, pos, node):
+            self._trie = trie
+            self._path = path
+            self._pos = pos
+            self._node = node
+
+        __bool__ = __nonzero__ = lambda self: True
+
+        @property
+        def is_set(self):
+            """Returns whether the node has value assigned to it."""
+            return self._node.value is not _SENTINEL
+
+        @property
+        def has_subtrie(self):
+            """Returns whether the node has any children."""
+            return bool(self._node.children)
+
+        def get(self, default=None):
+            """Returns node's value or the default if value is not assigned."""
+            v = self._node.value
+            return default if v is _SENTINEL else v
+
+        def set(self, value):
+            """Assigns value to the node."""
+            self._node.value = value
+
+        def setdefault(self, value):
+            """Assigns value to the node if one is not set then returns it."""
+            if self._node.value is _SENTINEL:
+                self._node.value = value
+            return self._node.value
+
+        @property
+        def key(self):
+            """Returns key of the node."""
+            if not hasattr(self, '_Step__key'):
+                self.__key = self._trie._key_from_path(self._path[:self._pos])
+            return self.__key
+
+        @property
+        def value(self):
+            """Returns node's value or raises KeyError."""
+            v = self._node.value
+            if v is _SENTINEL:
+                raise ShortKeyError(self.key)
+            return v
+
+    _NONE_STEP = _NoneStep()
+
+    def walk_towards(self, key):
+        """Yields nodes on the path to given node.
+
+        Args:
+            key: Key of the node to look for.
+
+        Yields:
+            :class:`pygtrie.Trie._Step` objects which can be used to extract or
+            set node's value as well as get node's key.
+
+            When representing nodes with assigned values, the objects can be
+            treated as ``(k, value)`` pairs denoting keys with associated values
+            encountered on the way towards the specified key.  This is
+            deprecated, prefer using ``key`` and ``value`` properties or ``get``
+            method of the object.
+
+        Raises:
+            KeyError: If node with given key does not exist.  It's all right if
+                they value is not assigned to the node provided it has a child
+                node.  Because the method is a generator, the exception is
+                raised only once a missing node is encountered.
+        """
+        node = self._root
+        path = self.__path_from_key(key)
+        pos = 0
+        while True:
+            yield self._Step(self, path, pos, node)
+            if pos == len(path):
+                break
+            node = node.children.get(path[pos])
+            if not node:
+                raise KeyError(key)
+            pos += 1
+
     def prefixes(self, key):
         """Walks towards the node specified by key and yields all found items.
 
@@ -896,30 +1013,27 @@ def prefixes(self, key):
             key: Key to look for.
 
         Yields:
-            ``(k, value)`` pairs denoting keys with associated values
-            encountered on the way towards the specified key.
+            :class:`pygtrie.Trie._Step` objects which can be used to extract or
+            set node's value as well as get node's key.
+
+            The objects can be treated as ``(k, value)`` pairs denoting keys
+            with associated values encountered on the way towards the specified
+            key.  This is deprecated, prefer using ``key`` and ``value``
+            properties of the object.
         """
-        node = self._root
-        path = self.__path_from_key(key)
-        pos = 0
-        while True:
-            if node.value is not _SENTINEL:
-                yield self._key_from_path(path[:pos]), node.value
-            if pos == len(path):
-                break
-            node = node.children.get(path[pos])
-            if not node:
-                break
-            pos += 1
+        try:
+            for step in self.walk_towards(key):
+                if step.is_set:
+                    yield step
+        except KeyError:
+            pass
 
     def shortest_prefix(self, key):
         """Finds the shortest prefix of a key with a value.
 
-        This is equivalent to taking the first object yielded by
-        :func:`Trie.prefixes` with a default of `(None, None)` if said method
-        yields no items.  As an added bonus, the pair in that case will be
-        a falsy value (as opposed to regular two-element tuple of ``None``
-        values).
+        This is roughly equivalent to taking the first object yielded by
+        :func:`Trie.prefixes` with additional handling for situations when no
+        prefixes are found.
 
         Example:
 
@@ -938,20 +1052,24 @@ def shortest_prefix(self, key):
             key: Key to look for.
 
         Returns:
-            ``(k, value)`` where ``k`` is the shortest prefix of ``key`` (it may
-            equal ``key``) and ``value`` is a value associated with that key.
-            If no node is found, ``(None, None)`` is returned.
+            :class:`pygtrie.Trie._Step` object (which can be used to extract or
+            set node's value as well as get node's key), or
+            a :class:`pygtrie.Trie._NoneStep` object (which is falsy value
+            simulating a _Step with ``None`` key and value) if no prefix is
+            found.
+
+            The object can be treated as ``(key, value)`` pair denoting key with
+            associated value of the prefix.  This is deprecated, prefer using
+            ``key`` and ``value`` properties of the object.
         """
-        return next(self.prefixes(key), _NONE_PAIR)
+        return next(self.prefixes(key), self._NONE_STEP)
 
     def longest_prefix(self, key):
         """Finds the longest prefix of a key with a value.
 
-        This is equivalent to taking the last object yielded by
-        :func:`Trie.prefixes` with a default of `(None, None)` if said method
-        yields no items.  As an added bonus, the pair in that case will be
-        a falsy value (as opposed to regular two-element tuple of ``None``
-        values).
+        This is roughly equivalent to taking the last object yielded by
+        :func:`Trie.prefixes` with additional handling for situations when no
+        prefixes are found.
 
         Example:
 
@@ -970,11 +1088,17 @@ def longest_prefix(self, key):
             key: Key to look for.
 
         Returns:
-            ``(k, value)`` where ``k`` is the longest prefix of ``key`` (it may
-            equal ``key``) and ``value`` is a value associated with that key.
-            If no node is found, ``(None, None)`` is returned.
+            :class:`pygtrie.Trie._Step` object (which can be used to extract or
+            set node's value as well as get node's key), or
+            a :class:`pygtrie.Trie._NoneStep` object (which is falsy value
+            simulating a _Step with ``None`` key and value) if no prefix is
+            found.
+
+            The object can be treated as ``(key, value)`` pair denoting key with
+            associated value of the prefix.  This is deprecated, prefer using
+            ``key`` and ``value`` properties of the object.
         """
-        ret = _NONE_PAIR
+        ret = self._NONE_STEP
         for ret in self.prefixes(key):
             pass
         return ret

test.py

@@ -286,7 +286,9 @@ def _do_test_find_one_prefix(self, trie_factory):
         none_pair = (None, None)
 
         def assert_pair(expected, got):
-            self.assertEqual(expected, got)
+            self.assertEqual(expected, tuple(got))
+            self.assertEqual(expected[0], got.key)
+            self.assertEqual(expected[1], got.value)
             if expected[0]:
                 self.assertTrue(got)
             else:
@@ -316,15 +318,74 @@ def _do_test_find_one_prefix(self, trie_factory):
         short_pair = (self.key_from_key(self._SHORT_KEY), 42)
         long_pair = (self.key_from_key(self._LONG_KEY), 42)
 
-        self.assertEqual([], list(t.prefixes(self._SHORT_PREFIXES[-1])))
-        self.assertEqual([], list(t.prefixes(self._OTHER_KEY)))
-        self.assertEqual([short_pair], list(t.prefixes(self._SHORT_KEY)))
-        self.assertEqual([short_pair],
-                         list(t.prefixes(self._LONG_PREFIXES[-1])))
-        self.assertEqual([short_pair, long_pair],
-                         list(t.prefixes(self._LONG_KEY)))
-        self.assertEqual([short_pair, long_pair],
-                         list(t.prefixes(self._VERY_LONG_KEY)))
+        def assert_prefixes(expected, *args):
+            got = list(t.prefixes(*args))
+            self.assertEqual(expected, [tuple(step) for step in got])
+            for e, g in zip(expected, got):
+                self.assertTrue(g)
+                self.assertEqual(e[0], g.key)
+                self.assertEqual(e[1], g.value)
+
+        assert_prefixes([], self._SHORT_PREFIXES[-1])
+        assert_prefixes([], self._OTHER_KEY)
+        assert_prefixes([short_pair], self._SHORT_KEY)
+        assert_prefixes([short_pair], self._LONG_PREFIXES[-1])
+        assert_prefixes([short_pair, long_pair], self._LONG_KEY)
+        assert_prefixes([short_pair, long_pair], self._VERY_LONG_KEY)
+
+    def _do_test_walk_towards(self, trie_factory):
+        """walk_towards method test."""
+        d = dict.fromkeys((self._SHORT_KEY, self._LONG_KEY), 42)
+        t = trie_factory(self._TRIE_CLS, d)
+
+        short_pair = (self.key_from_key(self._SHORT_KEY), 42)
+        long_pair = (self.key_from_key(self._LONG_KEY), 42)
+        none_pair = (None, None)
+
+        def assert_step(step):
+            self.assertTrue(step)
+            self.assertEqual(step[0], step.key)
+
+            is_set = step.key in (self.key_from_key(self._SHORT_KEY),
+                                  self.key_from_key(self._LONG_KEY))
+            self.assertEqual(is_set, step.is_set)
+            self.assertEqual(step.key != self.key_from_key(self._LONG_KEY),
+                             step.has_subtrie)
+
+            if is_set:
+                self.assertEqual(42, step.value)
+                self.assertEqual(42, step[1])
+                self.assertEqual(42, step.get('42'))
+                self.assertEqual(42, step.setdefault('42'))
+                self.assertEqual(42, step.value)
+                step.set('42')
+                self.assertEqual('42', step.value)
+                step.set(42)
+            else:
+                self.assertRaises(lambda: step.value)
+                self.assertRaises(lambda: step[1])
+                self.assertEqual('42', step.get('42'))
+                self.assertEqual(42, step.setdefault(42))
+                self.assertEqual(42, step.value)
+                step._node.value = pygtrie._SENTINEL
+
+        def assert_steps(key, raises=False):
+            try:
+                n = 0
+                for step in t.walk_towards(key):
+                    assert_step(step)
+                    n += 1
+                self.assertFalse(raises)
+                self.assertEqual(len(self.path_from_key(key)) + 1, n)
+            except KeyError:
+                self.assertTrue(raises)
+
+        assert_steps(self._SHORT_PREFIXES[-1])
+        assert_steps(self._OTHER_KEY[-1], raises=True)
+        assert_steps(self._SHORT_KEY)
+        assert_steps(self._LONG_PREFIXES[-1])
+        assert_steps(self._LONG_KEY)
+        assert_steps(self._VERY_LONG_KEY, raises=True)
 
     def _do_test_pickle(self, trie_factory):
         """https://github.com/google/pygtrie/issues/7"""