Repository: sympy/sympy

Issue:
`ask` function needs `except KeyError: ` condition
I am deeply interested in new assumption system, and trying with custom queries. But this led to an error when I meddled with assumptions.

I'll exemplify with `Q.mersenne`, which is provided as an example in sympy's documentation page.

```
>>> from sympy.assumptions import register_handler, ask, Q
>>> from sympy.assumptions.handlers import AskHandler
>>> class MersenneHandler(AskHandler):
...     # Mersenne numbers are in the form 2**n - 1, n integer
...     @staticmethod
...     def Integer(expr, assumptions):
...         from sympy import log
...         return ask(Q.integer(log(expr + 1, 2)))
>>> register_handler('mersenne', MersenneHandler)
>>> ask(Q.mersenne(7))
```

So far, so good.
But when I extended this to Symbol and taught it a silly fact: "If n is mersenne, then n is mersenne.", it raised an error.

```
>>> from sympy.assumptions import register_handler, ask, Q
>>> from sympy.assumptions.handlers import AskHandler
>>> from sympy.logic.boolalg import conjuncts
>>> class MersenneHandler(AskHandler):
...     # Mersenne numbers are in the form 2**n - 1, n integer
...     @staticmethod
...     def Integer(expr, assumptions):
...         from sympy import log
...         if ask(Q.integer(log(expr + 1, 2))):
...             return True
...     @staticmethod
...     def Symbol(expr, assumptions):
...         if expr in conjuncts(assumptions):
...             return True
>>> register_handler('mersenne', MersenneHandler)

>>> from sympy import Symbol
>>> n = Symbol('n', integer=True)
>>> ask(Q.mersenne(n), Q.mersenne(n))
```
This raised:
```
 Traceback (most recent call last)
   1300         # See if there's a straight-forward conclusion we can make for the inference
   1301         if local_facts_.is_Atom:
-> 1302             if key in known_facts_dict[local_facts_]:
   1303                 return True
   1304             if Not(key) in known_facts_dict[local_facts_]:

KeyError: Q.mersenne
```

(Of course, my original code wasn't about simply adding this kind of tautologic fact.)
I tested by adding `except KeyError: pass` condition and it resolved the error.

Maybe this small condition will be harmless to be added in the code?

Discussion:

Patch:
diff --git a/sympy/assumptions/ask.py b/sympy/assumptions/ask.py
index b699fdab87..2c071b7eaf 100644
--- a/sympy/assumptions/ask.py
+++ b/sympy/assumptions/ask.py
@@ -1295,32 +1295,22 @@ def ask(proposition, assumptions=True, context=global_assumptions):
         raise ValueError("inconsistent assumptions %s" % assumptions)
 
     if local_facts.clauses:
-        local_facts_ = CNF.CNF_to_cnf(local_facts)
 
-        # See if there's a straight-forward conclusion we can make for the inference
-        if local_facts_.is_Atom:
-            if key in known_facts_dict[local_facts_]:
-                return True
-            if Not(key) in known_facts_dict[local_facts_]:
-                return False
-        elif (isinstance(local_facts_, And) and
-              all(k in known_facts_dict for k in local_facts_.args)):
-            for assum in local_facts_.args:
-                if assum.is_Atom:
-                    if key in known_facts_dict[assum]:
-                        return True
-                    if Not(key) in known_facts_dict[assum]:
-                        return False
-                elif isinstance(assum, Not) and assum.args[0].is_Atom:
-                    if key in known_facts_dict[assum]:
-                        return False
-                    if Not(key) in known_facts_dict[assum]:
-                        return True
-        elif (isinstance(key, Predicate) and
-              isinstance(local_facts_, Not) and local_facts_.args[0].is_Atom):
-            if local_facts_.args[0] in known_facts_dict[key]:
+        if len(local_facts.clauses) == 1:
+            cl, = local_facts.clauses
+            f, = cl if len(cl)==1 else [None]
+            if f and f.is_Not and f.arg in known_facts_dict.get(key, []):
                 return False
 
+        for clause in local_facts.clauses:
+            if len(clause) == 1:
+                f, = clause
+                fdict = known_facts_dict.get(f.arg, None) if not f.is_Not else None
+                if fdict and key in fdict:
+                    return True
+                if fdict and Not(key) in known_facts_dict[f.arg]:
+                    return False
+
     # direct resolution method, no logic
     res = key(expr)._eval_ask(assumptions)
     if res is not None:


Test patch:
diff --git a/sympy/assumptions/tests/test_query.py b/sympy/assumptions/tests/test_query.py
index 3cd5752064..179996a3cd 100644
--- a/sympy/assumptions/tests/test_query.py
+++ b/sympy/assumptions/tests/test_query.py
@@ -2247,3 +2247,25 @@ def test_autosimp_used_to_fail():
     assert ask(Q.imaginary(0**(-I))) is False
     assert ask(Q.real(0**I)) is False
     assert ask(Q.real(0**(-I))) is False
+
+
+def test_custom_AskHandler():
+    from sympy.assumptions import register_handler, ask, Q
+    from sympy.assumptions.handlers import AskHandler
+    from sympy.logic.boolalg import conjuncts
+    from sympy import Symbol
+
+    class MersenneHandler(AskHandler):
+        @staticmethod
+        def Integer(expr, assumptions):
+            from sympy import log
+            if ask(Q.integer(log(expr + 1, 2))):
+                return True
+        @staticmethod
+        def Symbol(expr, assumptions):
+            if expr in conjuncts(assumptions):
+                return True
+    register_handler('mersenne', MersenneHandler)
+
+    n = Symbol('n', integer=True)
+    assert ask(Q.mersenne(n), Q.mersenne(n))