Improvements following feedback.

Author: haz

nnf/__init__.py

@@ -998,51 +998,51 @@ def pair(node: NNF) -> NNF:
 
         return pair(sentence)
 
-    def project(self: T_NNF, names: 't.FrozenSet[Name]') -> 'NNF':
-        """Dual of :meth:``forget``: will forget all variables not given"""
-        return self.forget(self.vars() - names)
+    def project(self, names: 't.Iterable[Name]') -> 'NNF':
+        """Dual of :meth:`forget: will forget all variables not given"""
+        return self.forget(self.vars() - frozenset(names))
 
-    def forget_aux(self: T_NNF) -> 'NNF':
+    def forget_aux(self) -> 'NNF':
         """Returns a theory that forgets all of the auxillary variables"""
-        aux_vars = frozenset({v for v in self.vars() if isinstance(v, Aux)})
-        return self.forget(aux_vars)
+        return self.forget(v for v in self.vars() if isinstance(v, Aux))
 
-    def forget(self: T_NNF, names: 't.FrozenSet[Name]') -> 'NNF':
+    def forget(self, names: 't.Iterable[Name]') -> 'NNF':
         """Forget a set of variables from the theory.
 
         Has the effect of returning a theory without the variables provided,
         such that every model of the new theory has an extension (i.e., an
         assignment) to the forgotten variables that is a model of the original
         theory.
 
-        :param names: A frozenset of the variable names to be forgotten
+        :param names: An iterable of the variable names to be forgotten
         """
 
         if self.decomposable():
             return self._forget_with_subs(names)
         else:
             return self._forget_with_shannon(names)
 
-    def _forget_with_subs(self: T_NNF, names: 't.FrozenSet[Name]') -> 'NNF':
+    def _forget_with_subs(self, names: 't.Iterable[Name]') -> 'NNF':
+
+        names = frozenset(names)
 
         @memoize
         def forget_recurse(node: NNF) -> NNF:
             if isinstance(node, Var) and node.name in names:
                 return true
             elif isinstance(node, Var):
                 return node
-            elif isinstance(node, And) or isinstance(node, Or):
-                new_children = map(forget_recurse, node.children)
-                return node.__class__(new_children)
+            elif isinstance(node, Internal):
+                return node.map(forget_recurse)
             else:
                 raise TypeError(node)
 
         return forget_recurse(self).simplify()
 
-    def _forget_with_shannon(self: T_NNF, vars: 't.FrozenSet[Name]') -> 'NNF':
+    def _forget_with_shannon(self, names: 't.Iterable[Name]') -> 'NNF':
         T = self
-        for v in vars:
-            T = t.cast(T_NNF, T.condition({v: True}) | T.condition({v: False}))
+        for v in frozenset(names) & self.vars():
+            T = T.condition({v: True}) | T.condition({v: False})
         return T.simplify()
 
     def deduplicate(self: T_NNF) -> T_NNF:

test_nnf.py

@@ -735,20 +735,21 @@ def test_iff(a: nnf.NNF, b: nnf.NNF):
 
 @given(NNF())
 def test_forget(sentence: nnf.NNF):
+    # Assumption to reduce the time in testing
+    assume(sentence.size() <= 15)
+
     # Test that forgetting a backbone variable doesn't change the theory
     T = sentence & Var('added_var')
     assert sentence.equivalent(T.forget({'added_var'}))
 
-    # Assumption to reduce the time in testing
-    assume(sentence.size() <= 15)
-
     # Test the tseitin projection
     assert sentence.equivalent(sentence.to_CNF().forget_aux())
 
     # Test that models of a projected theory are consistent with the original
-    assume(len(sentence.vars()) > 2)
-    vars = list(sentence.vars())[:2]
-    T = sentence.forget(vars)
+    names = list(sentence.vars())[:2]
+    T = sentence.forget(names)
+    assert not any([v in T.vars() for v in names])
+
     for m in T.models():
         assert sentence.condition(m).satisfiable()
 
@@ -757,8 +758,8 @@ def test_forget(sentence: nnf.NNF):
 def test_project(sentence: nnf.NNF):
     # Test that we get the same as projecting and forgetting
     assume(len(sentence.vars()) > 3)
-    vars1 = frozenset(list(sentence.vars())[:2])
-    vars2 = frozenset(list(sentence.vars())[2:])
+    vars1 = list(sentence.vars())[:2]
+    vars2 = list(sentence.vars())[2:]
     assert sentence.forget(vars1).equivalent(sentence.project(vars2))