Added functions from AntiORM

1 files changed, 44 insertions, 0 deletions

diff --git a/sqlparse/functions.py b/sqlparse/functions.py
new file mode 100644
index 0000000..693fe6a
--- /dev/null
+++ b/sqlparse/functions.py
@@ -0,0 +1,44 @@
+'''
+Created on 17/05/2012
+
+@author: piranna
+
+Several utility functions to extract info from the SQL sentences
+'''
+
+from sqlparse.filters  import ColumnsSelect, Limit
+from sqlparse.pipeline import Pipeline
+from sqlparse.tokens   import Keyword, Whitespace
+
+
+def GetLimit(stream):
+    """Function that return the LIMIT of a input SQL """
+    pipe = Pipeline()
+
+    pipe.append(Limit())
+
+    result = pipe(stream)
+    try:
+        return int(result)
+    except ValueError:
+        return result
+
+
+def GetColumns(stream):
+    """Function that return the colums of a SELECT query"""
+    pipe = Pipeline()
+
+    pipe.append(ColumnsSelect())
+
+    return pipe(stream)
+
+
+class IsType():
+    """Functor that return is the statement is of a specific type"""
+    def __init__(self, type):
+        self.type = type
+
+    def __call__(self, stream):
+        for token_type, value in stream:
+            if token_type in Whitespace: continue
+            return token_type in Keyword and value == self.type