Fix NameDicts so they keep track of the max_length properly and handle things like...

diff --git a/dns/namedict.py b/dns/namedict.py
index a9ae2b831f1044f46af555939b5809f9edcc51c9..03b5041fde82afa46780f47b5ce4a250835ecd25 100644 (file)
--- a/dns/namedict.py
+++ b/dns/namedict.py
@@ -1,4 +1,5 @@
 # Copyright (C) 2003-2007, 2009-2011 Nominum, Inc.
+# Copyright (C) 2016 Coresec Systems AB
 #
 # Permission to use, copy, modify, and distribute this software and its
 # documentation for any purpose with or without fee is hereby granted,
@@ -12,30 +13,69 @@
 # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT
 # OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+#
+# THE SOFTWARE IS PROVIDED "AS IS" AND CORESEC SYSTEMS AB DISCLAIMS ALL
+# WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL CORESEC
+# SYSTEMS AB BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR
+# CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
+# OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
+# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
+# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
 """DNS name dictionary"""
 
 import dns.name
+import collections
 
-class NameDict(dict):
+class NameDict(collections.MutableMapping):
 
     """A dictionary whose keys are dns.name.Name objects.
     @ivar max_depth: the maximum depth of the keys that have ever been
     added to the dictionary.
     @type max_depth: int
+    @ivar max_depth_items: the number of items of maximum depth
+    @type max_depth_items: int
     """
 
+    __slots__ = ["max_depth", "max_depth_items", "__store"]
+
     def __init__(self, *args, **kwargs):
-        super(NameDict, self).__init__(*args, **kwargs)
+        self.__store = dict()
         self.max_depth = 0
+        self.max_depth_items = 0
+        self.update(dict(*args, **kwargs))
+
+    def __update_max_depth(self,key):
+        if len(key) == self.max_depth:
+            self.max_depth_items = self.max_depth_items + 1
+        elif len(key) > self.max_depth:
+            self.max_depth = len(key)
+            self.max_depth_items = 1
+
+    def __getitem__(self, key):
+        return self.__store[key]
 
     def __setitem__(self, key, value):
         if not isinstance(key, dns.name.Name):
             raise ValueError('NameDict key must be a name')
-        depth = len(key)
-        if depth > self.max_depth:
-            self.max_depth = depth
-        super(NameDict, self).__setitem__(key, value)
+        self.__store[key] = value
+        self.__update_max_depth(key)
+
+    def __delitem__(self, key):
+        value = self.__store.pop(key)
+        if len(value) == self.max_depth:
+            self.max_depth_items = self.max_depth_items - 1
+        if self.max_depth_items == 0:
+            self.max_depth = 0
+            for k in self.__store:
+                self.__update_max_depth(k)
+
+    def __iter__(self):
+        return iter(self.__store)
+
+    def __len__(self):
+        return len(self.__store)
 
     def get_deepest_match(self, name):
         """Find the deepest match to I{name} in the dictionary.