add cnfvar helper for counting CNF trees

diff --git a/src/cnfvar.py b/src/cnfvar.py
index 1a2ee34..3dd7745 100644 (file)
--- a/src/cnfvar.py
+++ b/src/cnfvar.py
@@ -141,7 +141,7 @@ def walk_cnf(cnf, nested, fun, acc):
     """
     :type cnf: cnf list
     :type nested: bool
-    :type fun: ('a -> bool -> (cnf stuff) -> 'a)
+    :type fun: 'a -> bool -> (cnf stuff) -> 'a
     :type acc: 'a
     :rtype: 'a
     """
@@ -246,14 +246,21 @@ def is_cnf(root):
     well-formed member of the argument will cause the predicate to bail out
     with an exception during traversal.
     """
-    if isinstance (root, list):
-        cnf = root
-    else:
-        cnf = cnf_root (root)
+    cnf = cnf_root (root)
     if cnf is None:
         raise InvalidCNF (root)
     return walk_cnf(cnf, False, is_valid, {}) is not None
 
+
+def count_vars (root):
+    """
+    Traverse the cnf structure recursively, counting VAR objects (CNF lines).
+    """
+    cnf = cnf_root (root)
+    if cnf is None:
+        raise InvalidCNF (root)
+    return walk_cnf (cnf, True, lambda n, _nested, **_kwa: n + 1, 0)
+
 #
 #                               deserialization
 #
@@ -541,8 +548,7 @@ def renumber_vars(root, parent=None, toplevel=False):
     starting at a given offset without at the same time having that offset
     assigned as a parent.
     """
-    if isinstance(root, dict):
-        root = cnf_root (root)
+    root = cnf_root (root)
     i = parent or 0
     for var in root:
         i += 1
@@ -628,6 +634,8 @@ def cnf_root (root):
     containing the object itself, if it satisfies the criteria for a CNF_VAR,
     or the object contained in its toplevel field *cnf*.
     """
+    if isinstance (root, list):
+        return root
     if not isinstance(root, dict):
         raise TypeError(
             "Expected dictionary of CNF_VARs, got %s." % type(root))