From 35ed617fe41b184288892f733ce07d94391b76b1 Mon Sep 17 00:00:00 2001
From: Armin Ronacher <armin.ronacher@active-4.com>
Date: Mon, 17 May 2010 00:37:55 +0200
Subject: [PATCH] Added support for flashing categories.

This fixes #35.
---
 CHANGES              |  2 ++
 flask.py             | 33 +++++++++++++++++++++++++++++----
 tests/flask_tests.py | 29 +++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/CHANGES b/CHANGES
index ee8dd002..3a7d696d 100644
--- a/CHANGES
+++ b/CHANGES
@@ -8,6 +8,8 @@ Version 0.5
 
 Release date to be announced
 
+- added support for categories for flashed messages.
+
 Version 0.2
 -----------
 
diff --git a/flask.py b/flask.py
index 084b6593..d745027d 100644
--- a/flask.py
+++ b/flask.py
@@ -217,24 +217,49 @@ def get_template_attribute(template_name, attribute):
                    attribute)
 
 
-def flash(message):
+def flash(message, category='message'):
     """Flashes a message to the next request.  In order to remove the
     flashed message from the session and to display it to the user,
     the template has to call :func:`get_flashed_messages`.
 
+    .. versionchanged: 0.5
+       `category` parameter added.
+
     :param message: the message to be flashed.
+    :param category: the category for the message.  The following values
+                     are recommended: ``'message'`` for any kind of message,
+                     ``'error'`` for errors, ``'info'`` for information
+                     messages and ``'warning'`` for warnings.  However any
+                     kind of string and be used as category.
     """
-    session.setdefault('_flashes', []).append(message)
+    session.setdefault('_flashes', []).append((category, message))
 
 
-def get_flashed_messages():
+def get_flashed_messages(with_categories=False):
     """Pulls all flashed messages from the session and returns them.
     Further calls in the same request to the function will return
-    the same messages.
+    the same messages.  By default just the messages are returned,
+    but when `with_categories` is set to `True`, the return value will
+    be a list of tuples in the form ``(category, message)`` instead.
+
+    Example usage:
+
+    .. sourcecode:: html+jinja
+
+        {% for category, msg in get_flashed_messages(with_categories=true) %}
+          <p class=flash-{{ category }}>{{ msg }}
+        {% endfor %}
+
+    .. versionchanged:: 0.5
+       `with_categories` parameter added.
+
+    :param with_categories: set to `True` to also receive categories.
     """
     flashes = _request_ctx_stack.top.flashes
     if flashes is None:
         _request_ctx_stack.top.flashes = flashes = session.pop('_flashes', [])
+    if not with_categories:
+        return [x[1] for x in flashes]
     return flashes
 
 
diff --git a/tests/flask_tests.py b/tests/flask_tests.py
index c7cebb9c..f5dd12bb 100644
--- a/tests/flask_tests.py
+++ b/tests/flask_tests.py
@@ -157,6 +157,35 @@ class BasicFunctionalityTestCase(unittest.TestCase):
             assert flask.session.modified
             assert list(flask.get_flashed_messages()) == ['Zap', 'Zip']
 
+    def test_extended_flashing(self):
+        app = flask.Flask(__name__)
+        app.secret_key = 'testkey'
+
+        @app.route('/')
+        def index():
+            flask.flash(u'Hello World')
+            flask.flash(u'Hello World', 'error')
+            flask.flash(flask.Markup(u'<em>Testing</em>'), 'warning')
+            return ''
+
+        @app.route('/test')
+        def test():
+            messages = flask.get_flashed_messages(with_categories=True)
+            assert len(messages) == 3
+            assert messages[0] == ('message', u'Hello World')
+            assert messages[1] == ('error', u'Hello World')
+            assert messages[2] == ('warning', flask.Markup(u'<em>Testing</em>'))
+            return ''
+            messages = flask.get_flashed_messages()
+            assert len(messages) == 3
+            assert messages[0] == u'Hello World'
+            assert messages[1] == u'Hello World'
+            assert messages[2] == flask.Markup(u'<em>Testing</em>')
+
+        c = app.test_client()
+        c.get('/')
+        c.get('/test')
+
     def test_request_processing(self):
         app = flask.Flask(__name__)
         evts = []