From 5468b0a14e082c07c5db4f9d3c7f157c1df920a3 Mon Sep 17 00:00:00 2001 From: Victor Uriarte Date: Tue, 7 Jun 2016 21:13:33 -0700 Subject: Clean-up flatten upto --- sqlparse/filters/reindent.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) (limited to 'sqlparse') diff --git a/sqlparse/filters/reindent.py b/sqlparse/filters/reindent.py index 7d0e23a..8b3f7fa 100644 --- a/sqlparse/filters/reindent.py +++ b/sqlparse/filters/reindent.py @@ -22,19 +22,20 @@ class ReindentFilter(object): self._last_stmt = None def _flatten_up_to_token(self, token): - """Yields all tokens up to token plus the next one.""" - # helper for _get_offset - iterator = self._curr_stmt.flatten() - for t in iterator: + """Yields all tokens up to token but excluding current.""" + if token.is_group(): + token = next(token.flatten()) + + for t in self._curr_stmt.flatten(): yield t if t == token: raise StopIteration def _get_offset(self, token): raw = ''.join(map(text_type, self._flatten_up_to_token(token))) - line = raw.splitlines()[-1] + line = (raw or '\n').splitlines()[-1] # Now take current offset into account and return relative offset. - full_offset = len(line) - len(self.char * (self.width * self.indent)) + full_offset = len(line) - len(self.char * self.width * self.indent) return full_offset - self.offset def nl(self): -- cgit v1.2.1