Added memoize_generator and use it on IncludeStatement

2 files changed, 46 insertions, 4 deletions

diff --git a/sqlparse/filters.py b/sqlparse/filters.py
index 17dc960..92d7bd5 100644
--- a/sqlparse/filters.py
+++ b/sqlparse/filters.py
@@ -4,11 +4,11 @@ import re
 
 from os.path import abspath, join
 
-from sqlparse import sql
-from sqlparse import tokens as T
+from sqlparse        import sql, tokens as T
 from sqlparse.engine import FilterStack
-from sqlparse.tokens import (Comment, Comparison, Keyword, Name, Punctuation,
-                             String, Whitespace)
+from sqlparse.tokens import Comment, Comparison, Keyword, Name, Punctuation
+from sqlparse.tokens import String, Whitespace
+from sqlparse.utils  import memoize_generator
 
 
 # --------------------------
@@ -102,6 +102,7 @@ class IncludeStatement:
 
         self.detected = False
 
+    @memoize_generator
     def process(self, stack, stream):
         # Run over all tokens in the stream
         for token_type, value in stream:
diff --git a/sqlparse/utils.py b/sqlparse/utils.py
new file mode 100644
index 0000000..0dbb09f
--- /dev/null
+++ b/sqlparse/utils.py
@@ -0,0 +1,41 @@
+'''
+Created on 17/05/2012
+
+@author: piranna
+'''
+
+
+def memoize_generator(func):
+    """
+    Memoize decorator for generators
+
+    Store `func` results in a cache according to their arguments as 'memoize'
+    does but instead this works on decorators instead of regular functions.
+    Obviusly, this is only useful if the generator will always return the same
+    values for each specific parameters...
+    """
+    cache = {}
+
+    def wrapped_func(*args, **kwargs):
+        params = (args, kwargs)
+
+        # Look if cached
+        try:
+            cached = cache[params]
+
+        # Not cached, exec and store it
+        except KeyError:
+            cached = []
+
+            for item in func(*args, **kwargs):
+                cached.append(item)
+                yield item
+
+            cache[params] = cached
+
+        # Cached, yield its items
+        else:
+            for item in cached:
+                yield item
+
+    return wrapped_func