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

Code before:

    >>> 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), ~Q.integer(x))
    floor(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)), Q.complex(x) & Q.complex(y))
    ceiling(x) + ceiling(y)
    """
    arg = expr.args[0]
    if ask(Q.integer(arg), assumptions) or ask(Q.infinite(arg), assumptions):
        return arg

    if isinstance(arg, Add):
        simplified = []
        non_simplified = []
        for term in arg.args:
            if ask(Q.integer(term), assumptions) or isinstance(term, expr.func):
                simplified.append(term)
            else:
                non_simplified.append(term)
        return Add(*simplified) + expr.func(Add(*non_simplified))
    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,
}


Review comment (suggestion):
I think it would be important to come up with more descriptive variable names for these variables. `integer_terms` and `noninteger_terms` would be good except it's not quite accurate. Choosing good variables names can be hard - I struggle a lot with it sometimes - but it's very important. Don't about choosing long names as having short names matters much less than having descriptive names. 

Code after:

    >>> 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,
}