.

wcfs/internal/xbtree.py

@@ -164,18 +164,21 @@ class Bucket(object):
 #
 # The structure is represented as Tree and Bucket nodes.
 def StructureOf(znode):
+    ztype = _zclassify(znode)
+    """
     typ = type(znode)
     is_tree   = ("Tree"   in typ.__name__)
     is_set    = ("Set"    in typ.__name__)
     is_bucket = (("Bucket" in typ.__name__) or re.match("..Set", typ.__name__))
     is_map    = (not is_set)
+    """
 
-    if is_bucket:
-        keys, _ = zbcheck.crack_bucket(znode, is_map)
+    if ztype.is_zbucket:
+        keys, _ = zbcheck.crack_bucket(znode, ztype.is_map)
         return Bucket(*keys)
 
-    if is_tree:
-        kind, keys, children = zbcheck.crack_btree(znode, is_map)
+    if ztype.is_ztree:
+        kind, keys, children = zbcheck.crack_btree(znode, ztype.is_map)
         if kind == zbcheck.BTREE_EMPTY:
             return Tree([], Bucket())
 
@@ -195,7 +198,7 @@ def StructureOf(znode):
 # Restructure reorganizes ZODB BTree instance (not Tree) according to specified
 # structure.
 def Restructure(ztree, newStructure):
-    #assert istree(ztree)    # XXX reenable, -> use zbcheck.classify
+    #assert isztree(ztree)    # XXX reenable, -> use zbcheck.classify
     assert isinstance(newStructure, Tree)
 
     _bcheck(ztree)  # verify ztree before our tweaks
@@ -427,7 +430,12 @@ def _walkBFS(tree): # i[] of [](of nodes on each level)
 
 # _zwalkBFS, similarly to _walkBFS, walks ZODB BTree in breadth-first order layer by layer.
 def _zwalkBFS(ztree): # i[] of [](of nodes on each level)
-    # assert istree(ztree)
+    _ = _zclassify(ztree)
+    assert _.is_ztree
+
+    ztreeType    = type(ztree)
+    zbucketType  = ztree._bucket_type
+
     currentq = []
     nextq    = [ztree]
     while len(nextq) > 0:
@@ -436,9 +444,16 @@ def _zwalkBFS(ztree): # i[] of [](of nodes on each level)
         nextq    = []
         while len(currentq) > 0:
             znode = currentq.pop(0)
-            assert istree(znode) or isbucket(znode)
-            if istree(znode):
-                kind, keys, kids = zbcheck.crack_btree(znode)
+            ztype = _zclassify(znode)
+
+            assert ztype.is_ztree or ztype.is_zbucket
+            if ztype.is_zbucket:
+                assert type(znode) is zbucketType
+
+            if ztype.is_ztree:
+                assert type(znode) is ztreeType
+
+                kind, keys, kids = zbcheck.crack_btree(znode, ztype.is_map)
                 if kind == zbcheck.BTREE_EMPTY:
                     b = znode._bucket_type()
                     children = [b]
@@ -532,6 +547,31 @@ def _assertIncv(v):
         prev = v[i]
 
 
+# _zclassify returns kind of btree node znode is.
+# raises TypeError if znode is not a ZODB btree node.
+class _ZNodeType:
+    # .is_ztree     znode is a BTree node
+    # .is_zbucket   znode is a Bucket node
+    # .is_map       whether znode is k->v or just set(k)
+    pass
+def _zclassify(znode): # -> _ZNodeType
+    # XXX-> use zbcheck.classify ?
+    typ = type(znode)
+    is_ztree   = ("Tree"   in typ.__name__)
+    is_zset    = ("Set"    in typ.__name__)
+    is_zbucket = (("Bucket" in typ.__name__) or re.match("..Set", typ.__name__))
+    is_map     = (not is_zset)
+
+    if not (is_ztree or is_zbucket):
+        raise TypeError("type %r is not a ZODB BTree node" % typ)
+
+    _ = _ZNodeType()
+    _.is_ztree   = is_ztree
+    _.is_zbucket = is_zbucket
+    _.is_map     = is_map
+    return _
+
+
 # _bcheck performs full consistency checks on ztree provided by ZODB.
 #
 # The checks are what is provided by BTree.check and node._check().