Skip to content

Commit

Permalink
Fix Function.get_parameters(), add Funtion.get_window()
Browse files Browse the repository at this point in the history
  • Loading branch information
r33s3n6 authored and andialbrecht committed Apr 13, 2024
1 parent 617b8f6 commit e03b74e
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 1 deletion.
9 changes: 8 additions & 1 deletion sqlparse/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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."""
Expand Down
8 changes: 8 additions & 0 deletions tests/test_grouping.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down

0 comments on commit e03b74e

Please sign in to comment.