Repository: sympy/sympy
File: sympy/printing/smtlib.py (Python)

Code before:
                elif arity == 2:
                    p.symbol_table[func_head] = Callable[[float, float], float]
                else:
                    arg_types = [float] * arity
                    p.symbol_table[func_head] = Callable[arg_types, float]  # type: ignore[arg-type]

    # [FALLBACK RULES]
    for e in expr:
        for sym in e.atoms(Symbol, Function):
            if (
                sym.is_Symbol and
                sym not in p._known_constants and
                sym not in p.symbol_table
            ):
                log_warn(f"Could not infer type of `{sym}`. Defaulting to float.")
                p.symbol_table[sym] = float
            if (
                sym.is_Function and
                sym.func not in p._known_functions and
                sym.func not in p.symbol_table and  # <-- check head presence, not its type
                not sym.is_Piecewise
            ):
                raise TypeError(
                    f"Unknown type of undefined function `{sym}`. "
                    f"Must be mapped to ``str`` in known_functions or mapped to ``Callable[..]`` in symbol_table."
                )

    declarations = []
    if auto_declare:
        # Collect function heads to declare (both applied and unapplied)
        functions = {}
        for e in expr:
            for fnc in e.atoms(Function):
                if fnc.func not in p._known_functions and not fnc.is_Piecewise:
                    func_head = fnc.func
                    functions.setdefault(func_head.name, func_head)

        for raw in _raw_expr:
            if isinstance(raw, FunctionClass):
                functions.setdefault(raw.name, raw)
        # Gather constants from free symbols, excluding names that are function heads
        function_names = set(functions.keys())

        constants = {sym.name: sym for e in expr for sym in e.free_symbols
            if (sym not in p._known_constants) and (sym.name not in function_names)}

        declarations = [
            _auto_declare_smtlib(sym, p, log_warn)
            for sym in constants.values()
        ] + [
            generate_smtlib_declaration(fnc, p)

Review comment (suggestion):
Please avoid committing unrelated formatting changes. It makes it harder to review your pull request and these changes need to be removed for the code to be merged. When you `git commit` you can add the -p flag to go through each change one by one and make sure you actually want to commit it.

Code after:

    elif sym.is_Function:
        type_signature = p.symbol_table[type(sym)]
        assert callable(type_signature)
        type_signature = [p._known_types[_] for _ in type_signature.__args__]
        assert len(type_signature) > 0
        params_signature = f"({' '.join(type_signature[:-1])})"
        return_signature = type_signature[-1]
        return p._s_expr('declare-fun', [type(sym), params_signature, return_signature])

    else:
        log_warn(f"Non-Symbol/Function `{sym}` will not be declared.")
        return None


def _auto_assert_smtlib(e: Expr, p: SMTLibPrinter, log_warn: typing.Callable[[str], None]):
    if isinstance(e, Boolean) or (
        e in p.symbol_table and p.symbol_table[e] == bool
    ) or (
        e.is_Function and
        type(e) in p.symbol_table and
        p.symbol_table[type(e)].__args__[-1] == bool
    ):
        return p._s_expr('assert', [e])
    else:
        log_warn(f"Non-Boolean expression `{e}` will not be asserted. Converting to SMTLib verbatim.")
        return e


def _auto_infer_smtlib_types(
    *exprs: Basic,
    symbol_table: typing.Optional[dict] = None
) -> dict:
    # [TYPE INFERENCE RULES]
    # X is alone in an expr => X is bool
    # X in BooleanFunction.args => X is bool
    # X matches to a bool param of a symbol_table function => X is bool
    # X matches to an int param of a symbol_table function => X is int
    # X.is_integer => X is int
    # X == Y, where X is T => Y is T

    # [FALLBACK RULES]
    # see _auto_declare_smtlib(..)
    # X is not bool and X is not int and X is Symbol => X is float
    # else (e.g. X is Function) => error. must be specified explicitly.

    _symbols = dict(symbol_table) if symbol_table else {}

    def safe_update(syms: set, inf):
        for s in syms:
            assert s.is_Symbol