JSON exception handler for Flask

If you want Flask to spit json objects when exceptions are raised inside the application code, this exception handler can be useful:

from flask import jsonify
from werkzeug.exceptions import default_exceptions
from werkzeug.exceptions import HTTPException


class JSONExceptionHandler(object):

    def __init__(self, app=None):
        if app:
            self.init_app(app)

    def std_handler(self, error):
        response = jsonify(message=error.message)
        response.status_code = error.code if isinstance(error, HTTPException) else 500
        return response


    def init_app(self, app):
        self.app = app
        self.register(HTTPException)
        for code, v in default_exceptions.iteritems():
            self.register(code)

    def register(self, exception_or_code, handler=None):
        self.app.errorhandler(exception_or_code)(handler or self.std_handler)

To test it, try running the following app:

from exception_handler import *
from werkzeug.exceptions import default_exceptions, NotFound
from flask import Flask

app = Flask(__name__)
handler = JSONExceptionHandler(app)

class UserNotFound(NotFound):
    def __init__(self, uid):
        super(UserNotFound, self).__init__()
        self.message = "Couldn't found a user with uid={}.".format(uid)


@app.route("/<int:uid>")
def home(uid):
    if uid != 0:
        raise UserNotFound(uid)
    else:
        return jsonify({'user': "Foo", "uid": 0})


if __name__ == "__main__":
    app.run(debug=True)

If you run this app, a GET request at http://localhost:5000/0 will return a JSON:

{"user": "Foo", "uid": 0}

and a GET request at http://localhost:5000/1 will return a JSON:

{"message": "Couldn't found a user with uid=1."}

You could also create custom exception handlers:

def zero_division_handler(error):
    return jsonify({"message": "Don't divide by zero!!!!!1111"})

handler.register(ZeroDivisionError, handler=zero_division_handler)