From 2af1f0e29822610aae89b763a5935fa797e48f8c Mon Sep 17 00:00:00 2001 From: Zi-Xuan Fu Date: Wed, 27 Mar 2024 11:00:32 +0800 Subject: [PATCH] Fix Function.get_parameters(), add Funtion.get_window() --- sqlparse/sql.py | 9 ++++++++- tests/test_grouping.py | 8 ++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/sqlparse/sql.py b/sqlparse/sql.py index def06797..05e17748 100644 --- 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 index 0bf10c38..b39ff270 100644 --- 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():