updated DomainTrie to handle overlapping domains in the list

Author: jasonzhuyx

ml/utils/domain_trie.py

@@ -51,13 +51,17 @@ def _add_node(domain_pieces, tree, idx=0):
         if idx >= len(domain_pieces):
             tree['.'] = None
             return
-        # stop at "." because it is only building topper levels
+        part = domain_pieces[idx]
+        node = tree.get(part)
+        if node is None:
+            tree[part] = {}  # adding a new node
+        # NOTE: it is possible to have both "a.b.com" and "b.com" in the list,
+        #       and the order in domains list is not guaranteed.
         if "." not in tree:
-            part = domain_pieces[idx]
-            node = tree.get(part, {})
-            if node == {}:
-                tree[part] = node  # adding a new node
-            DomainTrie._add_node(domain_pieces, node, idx + 1)
+            DomainTrie._add_node(domain_pieces, tree[part], idx + 1)
+        else:
+            # only need to keep higher level domains
+            del tree[part]
 
     @staticmethod
     def _build_tree(domains):
@@ -68,7 +72,7 @@ def _build_tree(domains):
             if len(domain_pieces) > 1:
                 DomainTrie._add_node(domain_pieces[::-1], tree)
             else:
-                LOGGER.warn('Invalid domain: "%s"', domain)
+                LOGGER.warning('Invalid domain: "%s"', domain)
         return tree
 
     @staticmethod

tests/test_utils_domain_trie.py

@@ -14,6 +14,15 @@ def test_domain_trie(self):
             'xn--2342390.cn',
             'google.com',
             'test.google.com',
+            'abc.def.ghi.xyz.com',
+            'xyz.com',
+            'example.stuff'
+        ]
+        expected = [
+            '132.22.21.47',
+            'xn--2342390.cn',
+            'google.com',
+            'xyz.com',
             'example.stuff'
         ]
 
@@ -26,16 +35,17 @@ def test_domain_trie(self):
         self.assertIn('goo.example.stuff', trie)
         self.assertIn('google.com', trie)
         self.assertIn('test.google.com', trie)
+        self.assertIn('test.xyz.com', trie)
 
         self.assertNotIn('133.22.21.47', trie)
         self.assertNotIn('googl.com', trie)
         self.assertNotIn('google.api.com', trie)
         self.assertNotIn('google.cn', trie)
+        self.assertNotIn('xyz.ghi.com', trie)
         self.assertNotIn('stuff', trie)
 
-        domains.remove('test.google.com')
         result = ast.literal_eval(str(trie))
-        self.assertEqual(result, domains)
+        self.assertEqual(result, expected)
 
     def test_domain_trie_empty(self):
         domains = [