Repository: sympy/sympy
File: sympy/logic/algorithms/z3_wrapper.py (Python)

Code before:
from sympy.core.function import Function, UndefinedFunction, AppliedUndef
from sympy.functions.elementary.complexes import Abs
from sympy.functions.elementary.exponential import Pow
from sympy.functions.elementary.miscellaneous import Min, Max
from sympy.logic.boolalg import And, Or, Xor, Implies
from sympy.logic.boolalg import Not, ITE
from sympy.assumptions.relation.equality import StrictGreaterThanPredicate, StrictLessThanPredicate, GreaterThanPredicate, LessThanPredicate, EqualityPredicate
from sympy.external import import_module
from typing import Callable

# Define supported predicates as a constant set
supported_predicates = {
    Q.gt, Q.lt, Q.ge, Q.le, Q.ne, Q.eq, Q.positive, Q.negative, 
    Q.extended_negative, Q.extended_positive, Q.zero, Q.nonzero, 
    Q.nonnegative, Q.nonpositive, Q.extended_nonzero, 
    Q.extended_nonnegative, Q.extended_nonpositive
}

def is_supported_predicate(pred):
    if not isinstance(pred, AppliedPredicate):
        return False
    
    if pred.function in supported_predicates:
        return True

    return False

# def z3_satisfiable(expr, all_models=False):
#     if not isinstance(expr, EncodedCNF):
#         exprs = EncodedCNF()
#         exprs.add_prop(expr)
#         expr = exprs

#     z3 = import_module("z3")
#     if z3 is None:
#         raise ImportError("z3 is not installed")

#     s = encoded_cnf_to_z3_solver(expr, z3)

#     res = str(s.check())
#     if res == "unsat":
#         return False
#     elif res == "sat":
#         return z3_model_to_sympy_model(s.model(), expr)
#     else:
#         return None
    
def z3_satisfiable(expr, all_models=False):

    if not isinstance(expr, EncodedCNF):
        exprs = EncodedCNF()

Review comment (suggestion):
You could write this more concisely:
```python
return pred.function in supported_predicates
```

Code after:
from sympy.core.function import AppliedUndef
from sympy.functions.elementary.complexes import Abs
from sympy.functions.elementary.exponential import Pow
from sympy.functions.elementary.miscellaneous import Min, Max
from sympy.logic.boolalg import And, Or, Xor, Implies
from sympy.logic.boolalg import Not, ITE
from sympy.assumptions.relation.equality import StrictGreaterThanPredicate, StrictLessThanPredicate, GreaterThanPredicate, LessThanPredicate, EqualityPredicate
from sympy.external import import_module
from typing import Callable

supported_predicates = {
    Q.gt, Q.lt, Q.ge, Q.le, Q.ne, Q.eq, Q.positive, Q.negative,
    Q.extended_negative, Q.extended_positive, Q.zero, Q.nonzero,
    Q.nonnegative, Q.nonpositive, Q.extended_nonzero,
    Q.extended_nonnegative, Q.extended_nonpositive
}


def z3_satisfiable(expr, all_models=False):
    """
    Check satisfiability of a Boolean expression using the Z3 theorem prover.
    """
    if not isinstance(expr, EncodedCNF):
        exprs = EncodedCNF()
        exprs.add_prop(expr)
        expr = exprs

    z3 = import_module("z3")
    if z3 is None:
        raise ImportError("z3 is not installed")

    s = encoded_cnf_to_z3_solver(expr, z3)

    res = str(s.check())
    if res == "unsat":
        return False
    elif res == "sat":
        return z3_model_to_sympy_model(s.model(), expr)
    else:
        return None


def z3_model_to_sympy_model(z3_model, enc_cnf):
    rev_enc = {value: key for key, value in enc_cnf.encoding.items()}
    result = {}
    for var in z3_model:
        var_name = var.name()
        # Only process variables that start with 'd' followed by digits
        if var_name.startswith('d') and var_name[1:].isdigit():
            var_id = int(var_name[1:])
            if var_id in rev_enc: