Add merge method

Trie.merge method allows two tries to be merged structurally which is
often faster than looping over all items and setting them individually
in destination trie.  The operation doesn’t check that the tries are
of the same type so it may result in keys being transformed in one way
or another.

Issue: google#40

Author: mina86

.pylintrc

@@ -41,17 +41,19 @@ load-plugins=pylint.extensions.check_elif,
 # --enable=similarities". If you want to run only the classes checker, but have
 # no Warning level messages displayed, use"--disable=all --enable=classes
 # --disable=W"
-disable=locally-disabled,
+disable=consider-using-f-string,
+        locally-disabled,
         missing-return-type-doc,
         missing-type-doc,
         missing-yield-type-doc,
+        no-else-return,
         no-self-use,
+        super-with-arguments,
         too-few-public-methods,
         too-many-lines,
         too-many-public-methods,
-        useless-object-inheritance,
-        super-with-arguments,
-        consider-using-f-string
+        unidiomatic-typecheck,
+        useless-object-inheritance
 
 
 [REPORTS]

pygtrie.py

@@ -136,6 +136,16 @@ def add(self, parent, step):
     def require(self, parent, step):
         return self.node if self.step == step else self.add(parent, step)
 
+    def merge(self, other, queue):
+        """Moves children from other into this object."""
+        if type(other) == _OneChild and other.step == self.step:
+            queue.append((self.node, other.node))
+            return self
+        else:
+            children = _Children((self.step, self.node))
+            children.merge(other, queue)
+            return children
+
     def delete(self, parent, _step):
         parent.children = _EMPTY
 
@@ -172,6 +182,14 @@ def add(self, _parent, step):
     def require(self, _parent, step):
         return self.setdefault(step, _Node())
 
+    def merge(self, other, queue):
+        """Moves children from other into this object."""
+        for step, other_node in other.iteritems():
+            node = self.setdefault(step, other_node)
+            if node is not other_node:
+                queue.append((node, other_node))
+        return self
+
     def delete(self, parent, step):
         del self[step]
         if len(self) == 1:
@@ -196,6 +214,24 @@ def __init__(self):
         self.children = _EMPTY
         self.value = _EMPTY
 
+    def merge(self, other, overwrite):
+        """Move children from other node into this one.
+
+        Args:
+            other: Other node to move children and value from.
+            overwrite: Whether to overwrite existing node values.
+        """
+        queue = [(self, other)]
+        while queue:
+            lhs, rhs = queue.pop()
+            if lhs.value is _EMPTY or (overwrite and rhs.value is not _EMPTY):
+                lhs.value = rhs.value
+            if lhs.children is _EMPTY:
+                lhs.children = rhs.children
+            elif rhs.children is not _EMPTY:
+                lhs.children = lhs.children.merge(rhs.children, queue)
+            rhs.children = _EMPTY
+
     def iterate(self, path, shallow, iteritems):
         """Yields all the nodes with values associated to them in the trie.
 
@@ -484,6 +520,55 @@ def update(self, *args, **kwargs):  # pylint: disable=arguments-differ
             args = ()
         super(Trie, self).update(*args, **kwargs)
 
+    def merge(self, other, overwrite=False):
+        """Moves nodes from other trie into this one.
+
+        The merging happens at trie structure level and as such is different
+        than iterating over items of one trie and setting them in the other
+        trie.
+
+        The merging may happen between different types of tries resulting in
+        different (key, value) pairs in the destination trie compared to the
+        source.  For example, merging two :class:`pygtrie.StringTrie` objects
+        each using different separators will work as if the other trie had
+        separator of this trie.  Similarly, a :class:`pygtrie.CharTrie` may be
+        merged into a :class:`pygtrie.StringTrie` but when keys are read those
+        will be joined by the separator.  For example:
+
+            >>> import pygtrie
+            >>> st = pygtrie.StringTrie(separator='.')
+            >>> st.merge(pygtrie.StringTrie({'foo/bar': 42}))
+            >>> list(st.items())
+            [('foo.bar', 42)]
+            >>> st.merge(pygtrie.CharTrie({'baz': 24}))
+            >>> sorted(st.items())
+            [('b.a.z', 24), ('foo.bar', 42)]
+
+        Not all tries can be merged into other tries.  For example,
+        a :class:`pygtrie.StringTrie` may not be merged into
+        a :class:`pygtrie.CharTrie` because the latter imposes a requirement for
+        each component in the key to be exactly one character while in the
+        former components may be arbitrary length.
+
+        Note that the other trie is cleared and any references or iterators over
+        it are invalidated.  To preserve other’s value it needs to be copied
+        first.
+
+        Args:
+            other: Other trie to move nodes from.
+            overwrite: Whether to overwrite existing values in this trie.
+        """
+        if isinstance(self, type(other)):
+            self._merge_impl(self, other, overwrite=overwrite)
+        else:
+            other._merge_impl(self, other, overwrite=overwrite) # pylint: disable=protected-access
+        other.clear()
+
+    @classmethod
+    def _merge_impl(cls, dst, src, overwrite):
+        # pylint: disable=protected-access
+        dst._root.merge(src._root, overwrite=overwrite)
+
     def copy(self, __make_copy=lambda x: x):
         """Returns a shallow copy of the object."""
         # pylint: disable=protected-access
@@ -1637,6 +1722,13 @@ def fromkeys(cls, keys, value=None, separator='/'):  # pylint: disable=arguments
             trie[key] = value
         return trie
 
+    @classmethod
+    def _merge_impl(cls, dst, src, overwrite):
+        if not isinstance(dst, StringTrie):
+            raise TypeError('%s cannot be merged into a %s' % (
+                type(src).__name__, type(dst).__name__))
+        super(StringTrie, cls)._merge_impl(dst, src, overwrite=overwrite)
+
     def __str__(self):
         if not self:
             return '%s(separator=%s)' % (type(self).__name__, self._separator)

test.py

@@ -1114,5 +1114,43 @@ def __len__(self):
             self.assertNotEqual(trie, dictionary)
 
 
+class MergeTest(unittest.TestCase):
+    """Tests for merge method."""
+
+    def test_merge_tries(self):
+        trie = pygtrie.Trie({'foo': 1, 'bar': 2})
+
+        def test(want, src, overwrite=False):
+            trie.merge(src, overwrite=overwrite)
+            self.assertEqual(0, len(src))
+            self.assertEqual(pygtrie.Trie(want), trie)
+
+        test({'foo': 1, 'bar': 2, 'baz': 3},
+             pygtrie.Trie({'bar': 0, 'baz': 3}))
+        test({'foo': 1, 'bar': 2, 'baz': 4, 'fo': 5},
+             pygtrie.Trie({'baz': 4, 'fo': 5}), overwrite=True)
+        test({'foo': 1, 'bar': 2, 'baz': 4, 'fo': 5, 'qux': 6},
+             pygtrie.CharTrie({'qux': 6}))
+
+        st = pygtrie.StringTrie({'foo/bar/baz': 42})
+        self.assertRaises(TypeError, trie.merge, st)
+
+    def test_merge_string_tries(self):
+
+        def test(want, other, overwrite):
+            trie = pygtrie.StringTrie({'foo/bar': 42})
+            trie.merge(other, overwrite=overwrite)
+            self.assertEqual(want, dict(trie.items()))
+
+        test({'foo/bar': 42, 'bar/baz': 2},
+             pygtrie.StringTrie({'foo.bar': 4, 'bar.baz': 2}, separator='.'),
+             False)
+        test({'foo/bar': 4, 'bar/baz': 2},
+             pygtrie.StringTrie({'foo.bar': 4, 'bar.baz': 2}, separator='.'),
+             True)
+        test({'foo/bar': 42, 'q/u/x': 2}, pygtrie.Trie({'qux': 2}), False)
+        test({'foo/bar': 42, 'q/u/x': 2}, pygtrie.CharTrie({'qux': 2}), False)
+
+
 if __name__ == '__main__':
     unittest.main()

version-history.rst

@@ -1,7 +1,9 @@
 Version History
 ---------------
 
-2.4.3: TBD
+2.5.0: TBD
+
+- Add :func:`pygtrie.Trie.merge` which merges structures of two tries.
 
 - Fix :func:`pygtrie.Trie.__eq__` implementation such that key values
   are taken into consideration rather than just looking at trie