Search code examples
pythonoverridingstatementsuser-defined

Is it possible to declare user-defined statements in python (same syntax of assert)?


For example, what assert does:

assert <boolean assertion to check>, <error string/value outputted if that assertion False>

What I would like to do:

mycommand <boolean assertion to check>, <value that will get returned if assertion is False>

It would be the same as doing:

if not <boolean assertion to check>: return <value>

but concisely, with a simple one-liner. I know I should just use this last one, but I am more interested on finding out how much you can override/customize python.


Solution

  • You can not use any "user-defined" syntax, but you could just use assert itself and use a decorator to make the function return the value instead of raising an AssertionError:

    def assert_to_return(f):
        def _f(*args, **kwargs):
            try:
                return f(*args, **kwargs)
            except AssertionError as e:
                return e.args[0] if e.args else None
        return _f
    
    @assert_to_return
    def div(x, y):
        assert y != 0, float("nan")
        return x / y
    
    print(div(42, 5)) # 8.4
    print(div(42, 0)) # nan
    

    This way, the assert <condition>, <value> line means "check whether <condition> holds, otherwise immediately return <value>", i.e. pretty much exactly what you wanted. If the <value> is omitted, the function will just return None in that case, as defined in the decorator. Also note that <value> is only evaluated if the assert is violated, i.e. it could also contain an expensive function call or side-effects that you do not want to occur otherwise.

    Of course, while other exceptions will work just fine, any AssertionError raised by assert statements in functions your function is calling will also bei caught and handled by the decorator. If this is a problem, you could either (a) wrap the <value> in some special wrapper object and only return those in _f, otherwise re-raise the AssertionError, or (b) wrap all the functions you call in your function that might also raise an AssertionError so that their AssertionErrors are wrapped in another exception.