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
})
