Making Python's `assert` throw an exception that I choose

Can I make assert throw an exception that I choose instead of AssertionError?

UPDATE:

I'll explain my motivation: Up to now, I've had assertion-style tests that raised my own exceptions; For example, when you created a Node object with certain arguments, it would check if the arguments were good for creating a node, and if not it would raise NodeError.

But I know that Python has a -o mode in which asser开发者_如何学运维ts are skipped, which I would like to have available because it would make my program faster. But I would still like to have my own exceptions. That's why I want to use assert with my own exceptions.

This will work. But it's kind of crazy.

try:
    assert False, "A Message"
except AssertionError, e:
    raise Exception( e.args )

Why not the following? This is less crazy.

if not someAssertion: raise Exception( "Some Message" )

It's only a little wordier than the assert statement, but doesn't violate our expectation that assert failures raise AssertionError.

Consider this.

def myAssert( condition, action ):
    if not condition: raise action

Then you can more-or-less replace your existing assertions with something like this.

myAssert( {{ the original condition }}, MyException( {{ the original message }} ) )

Once you've done this, you are now free to fuss around with enable or disabling or whatever it is you're trying to do.

Also, read up on the warnings module. This may be exactly what you're trying to do.

How about this?

>>> def myraise(e): raise e
... 
>>> cond=False
>>> assert cond or myraise(RuntimeError)
Traceback (most recent call last):
  File "", line 1, in 
  File "", line 1, in myraise
RuntimeError

Never use an assertion for logic! Only for optional testing checks. Remember, if Python is running with optimizations turned on, asserts aren't even compiled into the bytecode. If you're doing this, you obviously care about the exception being raised and if you care, then you're using asserts wrong in the first place.

Python also skips if __debug__: blocks when run with -o option. The following code is more verbose, but does what you need without hacks:

def my_assert(condition, message=None):
    if not condition:
        raise MyAssertError(message)

if __debug__: my_assert(condition, message)

You can make it shorter by moving if __debug__: condition inside my_assert(), but then it will be called (without any action inside) when optimization is enabled.

You can let a context manager do the conversion for you, inside a with block (which may contain more than one assertion, or more code and function calls or what you want.

from __future__ import with_statement
import contextlib

@contextlib.contextmanager
def myassert(exctype):
    try:
        yield
    except AssertionError, exc:
        raise exctype(*exc.args)

with myassert(ValueError):
    assert 0, "Zero is bad for you"

See a previous version of this answer for substituting constructed exception objects directly (KeyError("bad key")), instead of reusing the assertions' argument(s).

In Python 2.6.3 at least, this will also work:

class MyAssertionError (Exception):
    pass

AssertionError = MyAssertionError

assert False, "False"

Traceback (most recent call last):
  File "assert.py", line 8, in <module>
    assert False, "False"
__main__.MyAssertionError: False

If you want to use asserts for this, this seems to work pretty well:

>>> def raise_(e): raise e
...
>>> x = -2
>>> assert x >= 0, raise_(ValueError('oops'))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 1, in raise_
ValueError: oops

Note that in the assert, the stuff after the comma will only be evaluated if the condition is false, so the ValueError will only be created and raised when needed.

To see if try has any overhead I tried this experiment

here is myassert.py

def myassert(e):
    raise e


def f1():                       #this is the control for the experiment
    cond=True

def f2():
    cond=True
    try:
        assert cond, "Message"
    except AssertionError, e:
        raise Exception(e.args)

def f3():
    cond=True
    assert cond or myassert(RuntimeError)

def f4():
    cond=True
    if __debug__:
        raise(RuntimeError)

$ python -O -mtimeit -n100 -r1000 -s'import myassert' 'myassert.f1()'
100 loops, best of 1000: 0.42 usec per loop
$ python -O -mtimeit -n100 -r1000 -s'import myassert' 'myassert.f2()'
100 loops, best of 1000: 0.479 usec per loop
$ python -O -mtimeit -n100 -r1000 -s'import myassert' 'myassert.f3()'
100 loops, best of 1000: 0.42 usec per loop
$ python -O -mtimeit -n100 -r1000 -s'import myassert' 'myassert.f4()'
100 loops, best of 1000: 0.42 usec per loop