Fix Function.get_parameters(), add Funtion.get_window()

andialbrecht · Apr 13, 2024 · e03b74e · e03b74e
1 parent 617b8f6
commit e03b74e
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/sqlparse/sql.py b/sqlparse/sql.py
@@ -623,7 +623,7 @@ class Function(NameAliasMixin, TokenList):
 
     def get_parameters(self):
         """Return a list of parameters."""
-        parenthesis = self.tokens[-1]
+        parenthesis = self.token_next_by(i=Parenthesis)[1]
         result = []
         for token in parenthesis.tokens:
             if isinstance(token, IdentifierList):
@@ -633,6 +633,13 @@ def get_parameters(self):
                 result.append(token)
         return result
 
+    def get_window(self):
+        """Return the window if it exists."""
+        over_clause = self.token_next_by(i=Over)
+        if not over_clause:
+            return None
+        return over_clause[1].tokens[-1]
+
 
 class Begin(TokenList):
     """A BEGIN/END block."""

diff --git a/tests/test_grouping.py b/tests/test_grouping.py
@@ -392,6 +392,14 @@ def test_grouping_function():
     p = sqlparse.parse('foo(null, bar)')[0]
     assert isinstance(p.tokens[0], sql.Function)
     assert len(list(p.tokens[0].get_parameters())) == 2
+    p = sqlparse.parse('foo(5) over win1')[0]
+    assert isinstance(p.tokens[0], sql.Function)
+    assert len(list(p.tokens[0].get_parameters())) == 1
+    assert isinstance(p.tokens[0].get_window(), sql.Identifier)
+    p = sqlparse.parse('foo(5) over (PARTITION BY c1)')[0]
+    assert isinstance(p.tokens[0], sql.Function)
+    assert len(list(p.tokens[0].get_parameters())) == 1
+    assert isinstance(p.tokens[0].get_window(), sql.Parenthesis)
 
 
 def test_grouping_function_not_in():