From c86449d39cde438e153055e0f7731f0d15b62964 Mon Sep 17 00:00:00 2001 From: Oleg Broytman Date: Sun, 14 Aug 2016 14:59:11 +0300 Subject: [PATCH] Simplify get_tokens_str; rename get_tokens_str -> tlist2str Use .flatten() and avoid StringIO - make the output string directly. --- mysql2sql/print_tokens.py | 13 ++----------- tests/test_tokens.py | 4 ++-- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/mysql2sql/print_tokens.py b/mysql2sql/print_tokens.py index 2ed6341..142391f 100644 --- a/mysql2sql/print_tokens.py +++ b/mysql2sql/print_tokens.py @@ -1,12 +1,5 @@ import sys -try: - from cStringIO import StringIO -except ImportError: - try: - from StringIO import StringIO - except ImportError: - from io import StringIO def print_tokens(token_list, outfile=sys.stdout): @@ -14,7 +7,5 @@ def print_tokens(token_list, outfile=sys.stdout): outfile.write(token.normalized) -def get_tokens_str(token_list): - sio = StringIO() - print_tokens(token_list, outfile=sio) - return sio.getvalue() +def tlist2str(token_list): + return ''.join(token.normalized for token in token_list.flatten()) diff --git a/tests/test_tokens.py b/tests/test_tokens.py index 1746f7b..c39cd31 100755 --- a/tests/test_tokens.py +++ b/tests/test_tokens.py @@ -5,7 +5,7 @@ import unittest from sqlparse import parse from mysql2sql.process_tokens import requote_names -from mysql2sql.print_tokens import get_tokens_str +from mysql2sql.print_tokens import tlist2str from tests import main @@ -13,7 +13,7 @@ class TestTokens(unittest.TestCase): def test_requote(self): parsed = parse("select * from `T`")[0] requote_names(parsed) - query = get_tokens_str(parsed) + query = tlist2str(parsed) self.assertEqual(query, 'SELECT * FROM "T"') -- 2.39.5