From b052e09cd6c7b0ddcfe42eab61f8ea58224d83da Mon Sep 17 00:00:00 2001
From: Matt Wright <mdw1980@gmail.com>
Date: Thu, 23 Aug 2012 20:47:48 -0400
Subject: [PATCH] Polish

---
 flask_security/core.py | 28 +++++++++-------------------
 1 file changed, 9 insertions(+), 19 deletions(-)

diff --git a/flask_security/core.py b/flask_security/core.py
index 7fffed3..5f042a5 100644
--- a/flask_security/core.py
+++ b/flask_security/core.py
@@ -158,11 +158,12 @@ def _get_state(app, datastore, **kwargs):
         login_manager=_get_login_manager(app),
         principal=_get_principal(app),
         pwd_context=_get_pwd_context(app),
-        context_processors={},
         remember_token_serializer=_get_serializer(app, 'remember'),
         login_serializer=_get_serializer(app, 'login'),
         reset_serializer=_get_serializer(app, 'reset'),
-        confirm_serializer=_get_serializer(app, 'confirm')
+        confirm_serializer=_get_serializer(app, 'confirm'),
+        _context_processors={},
+        _send_mail_task=None
     ))
 
     return _SecurityState(**kwargs)
@@ -217,31 +218,20 @@ class _SecurityState(object):
     def __init__(self, **kwargs):
         for key, value in kwargs.items():
             setattr(self, key.lower(), value)
-        self._send_mail_task = None
 
     def _add_ctx_processor(self, endpoint, fn):
-        c = self.context_processors
-
-        if endpoint not in c:
-            c[endpoint] = []
-
-        if fn not in c[endpoint]:
-            c[endpoint].append(fn)
+        group = self._context_processors.setdefault(endpoint, [])
+        fn not in group and group.append(fn)
 
     def _run_ctx_processor(self, endpoint):
         rv, fns = {}, []
-
-        for g in ['all', endpoint]:
-            if g in self.context_processors:
-                fns += self.context_processors[g]
-
-        for fn in fns:
-            rv.update(fn())
-
+        for g in [None, endpoint]:
+            for fn in self._context_processors.setdefault(g, []):
+                rv.update(fn())
         return rv
 
     def context_processor(self, fn):
-        self._add_ctx_processor('all', fn)
+        self._add_ctx_processor(None, fn)
 
     def forgot_password_context_processor(self, fn):
         self._add_ctx_processor('forgot_password', fn)