def _reformat_helper(input_code, imports): from pyflyby._imports2s import reformat_import_statements if PYFLYBY_START_MSG in input_code: before, bmarker, middle = input_code.partition(PYFLYBY_START_MSG) else: before, bmarker, middle = "", "", input_code if PYFLYBY_END_MSG in middle: middle, emarker, after = middle.partition(PYFLYBY_END_MSG) else: middle, emarker, after = middle, "", "" if imports is not None: transform = SourceToSourceFileImportsTransformation(middle) if isinstance(imports, str): imports = [imports] for imp in imports: assert isinstance(imp, str) if not imp.strip(): continue transform.add_import(Import(imp)) middle = str(transform.output()) return reformat_import_statements(before + bmarker + middle + emarker + after)
def test_reformat_import_statements_multi_star_1(): input = PythonBlock(dedent(''' from mod1 import * from mod2 import * ''').lstrip()) output = reformat_import_statements(input) expected = PythonBlock(dedent(''' from mod1 import * from mod2 import * ''').lstrip()) assert output == expected
def test_reformat_continuation_comments_1(): input = dedent(''' import b , a # x\\ # y ''').lstrip() output = str(reformat_import_statements(input)) expected = dedent(''' import a import b # x\\ # y ''').lstrip() assert expected == output
def test_reformat_import_statements_shadowed_1(): input = PythonBlock(dedent(''' import a, a2 as a, b, b2 as b, b, c2 as c, d2 as d from d import d import c ''').lstrip()) output = reformat_import_statements(input) expected = PythonBlock(dedent(''' import a2 as a import b import c from d import d ''').lstrip()) assert output == expected
def test_reformat_import_statements_1(): input = PythonBlock(dedent(''' from foo import bar2 as bar2x, bar1 import foo.bar3 as bar3x import foo.bar4 import foo.bar0 as bar0 ''').lstrip(), filename="/foo/test_reformat_import_statements_1.py") output = reformat_import_statements(input) expected = PythonBlock(dedent(''' import foo.bar4 from foo import bar1, bar2 as bar2x, bar3 as bar3x from foo import bar0 ''').lstrip(), filename="/foo/test_reformat_import_statements_1.py") assert output == expected
def test_reformat_import_statements_star_1(): input = PythonBlock(dedent(''' from mod1 import f1b from mod1 import f1a from mod2 import * from mod2 import f2b as F2B from mod2 import f2a as F2A from mod3 import f3b from mod3 import f3a ''').lstrip(), filename="/foo/test_reformat_import_statements_star_1.py") output = reformat_import_statements(input) expected = PythonBlock(dedent(''' from mod1 import f1a, f1b from mod2 import * from mod2 import f2a as F2A, f2b as F2B from mod3 import f3a, f3b ''').lstrip(), filename="/foo/test_reformat_import_statements_star_1.py") assert output == expected
def _recv(msg): data = msg["content"]["data"] if data["type"] == FORMATTING_IMPORTS: imports = data.get('imports', None) input_code = data["input_code"] if imports is not None: transform = SourceToSourceFileImportsTransformation(input_code) if isinstance(imports, str): imports = [imports] for imp in imports: assert isinstance(imp, str) if not imp.strip(): continue transform.add_import(Import(imp)) input_code = str(transform.output()) fmt_code = reformat_import_statements(input_code) comm.send({ "formatted_code": str(fmt_code), "type": FORMATTING_IMPORTS })