Repository: sympy/sympy
File: sympy/assumptions/refine.py (Python)

Code before:

    Examples
    ========

    >>> from sympy import Symbol, refine, Q, floor
    >>> x = Symbol('x')
    >>> y = Symbol('y')
    >>> refine(floor(x), Q.integer(x))
    x
    >>> refine(floor(x + y), Q.integer(x))
    x + floor(y)
    """
    from sympy.functions.elementary.integers import floor
    arg = expr.args[0]
    if ask(Q.integer(arg), assumptions):
        return arg

    if isinstance(arg, Add):
        if ask(Q.integer(arg.args[0]), assumptions):
            return arg.args[0] + floor(arg.args[1])
        elif ask(Q.integer(expr.args[1]), assumptions):
            return arg.args[1] + floor(arg.args[0])
    return expr


def refine_ceiling(expr, assumptions):
    """
    handler for ceiling function

    Examples
    ========

    >>> from sympy import Symbol, refine, Q, ceiling
    >>> x = Symbol('x')
    >>> refine(ceiling(x), Q.integer(x))
    x
    >>> y = Symbol('y')
    >>> refine(ceiling(x + y), Q.integer(x))
        x + ceiling(y)
    """
    from sympy.functions.elementary.integers import ceiling
    arg = expr.args[0]
    if ask(Q.integer(arg), assumptions):
        return arg

    if isinstance(arg, Add):
        if ask(Q.integer(arg.args[0]), assumptions):
            return arg.args[0] + ceiling(arg.args[1])
        elif ask(Q.integer(arg.args[1]), assumptions):
            return arg.args[1] + ceiling(arg.args[0])


Review comment (refactor):
I think it would make more sense to have a single `refine_floor_ceiling` function to avoid duplicate code. Unless you think that there would be situations where you would only want to refine either floor or ceiling?

Code after:

    Examples
    ========

    >>> from sympy import Symbol, refine, Q, floor, ceiling
    >>> x = Symbol('x')
    >>> y = Symbol('y')
    >>> refine(floor(x), Q.integer(x))
    x
    >>> refine(ceiling(x), Q.integer(x))
    x
    >>> refine(floor(x + y), Q.integer(x))
    x + floor(y)
    >>> refine(ceiling(x + y), Q.integer(x))
    x + ceiling(y)
    >>> refine(ceiling(x), Q.infinite(x))
    x
    >>> refine(ceiling(ceiling(x) + ceiling(y)))
    ceiling(x) + ceiling(y)
    """
    from sympy.functions.elementary.integers import floor, ceiling
    arg = expr.args[0]
    if ask(Q.integer(arg), assumptions) or ask(Q.infinite(arg), assumptions):
        return arg

    if isinstance(arg, Add):
        gaussian_integer_terms = []
        nongausian_intergers_terms = []
        for term in arg.args:
            if ask(Q.integer(term), assumptions) or isinstance(term, (floor, ceiling)):
                gaussian_integer_terms.append(term)
            else:
                nongausian_intergers_terms.append(term)
        return Add(*gaussian_integer_terms) + expr.func(Add(*nongausian_intergers_terms))
    return expr


handlers_dict: dict[str, Callable[[Basic, Boolean | bool], Expr]] = {
    'Abs': refine_abs,
    'Pow': refine_Pow,
    'atan2': refine_atan2,
    're': refine_re,
    'im': refine_im,
    'arg': refine_arg,
    'sign': refine_sign,
    'MatrixElement': refine_matrixelement,
    'cos': refine_sin_cos,
    'sin': refine_sin_cos,
    'floor': refine_floor_ceiling,
    'ceiling' : refine_floor_ceiling,
}