def test_add_existing_import_normal_import_aliased(self):
tree = ast.parse('import a.b.c as d')
self.assertEqual(
'a.b', import_utils.add_import(tree, 'a.b', from_import=False))
self.assertEqual(
'd', import_utils.add_import(tree, 'a.b.c', from_import=False))
self.assertEqual('import a.b\nimport a.b.c as d\n', pasta.dump(tree))
def test_add_import_with_asname_with_conflict(self):
tree = ast.parse('def c(): pass\n')
self.assertEqual(
'c_1',
import_utils.add_import(tree, 'a.b', asname='c', from_import=True))
self.assertEqual('from a import b as c_1\ndef c():\n pass\n',
pasta.dump(tree))
def test_merge_from_import(self):
tree = ast.parse('from a.b import c')
# x is explicitly not merged
self.assertEqual(
'x',
import_utils.add_import(tree, 'a.b.x', merge_from_imports=False))
self.assertEqual('from a.b import x\nfrom a.b import c\n',
pasta.dump(tree))
# y is allowed to be merged and is grouped into the first matching import
self.assertEqual(
'y', import_utils.add_import(tree,
'a.b.y',
merge_from_imports=True))
self.assertEqual('from a.b import x, y\nfrom a.b import c\n',
pasta.dump(tree))
def test_add_single_name_from_import_with_asname(self):
tree = ast.parse('')
self.assertEqual(
'bar',
import_utils.add_import(tree,
'foo',
asname='bar',
from_import=True))
self.assertEqual('import foo as bar\n', pasta.dump(tree))
def test_add_from_import_with_asname(self):
tree = ast.parse('')
self.assertEqual(
'd',
import_utils.add_import(tree,
'a.b.c',
asname='d',
from_import=True))
self.assertEqual('from a.b import c as d\n', pasta.dump(tree))
def test_add_normal_import_with_asname(self):
tree = ast.parse('')
self.assertEqual(
'd',
import_utils.add_import(tree,
'a.b.c',
asname='d',
from_import=False))
self.assertEqual('import a.b.c as d\n', pasta.dump(tree))
def test_add_import_after_docstring(self):
tree = ast.parse('\'Docstring.\'')
self.assertEqual('a', import_utils.add_import(tree, 'a'))
self.assertEqual('\'Docstring.\'\nimport a\n', pasta.dump(tree))
def test_add_existing_import_aliased_with_asname(self):
tree = ast.parse('from a.b import c as d')
self.assertEqual('d', import_utils.add_import(tree,
'a.b.c',
asname='e'))
self.assertEqual('from a.b import c as d\n', pasta.dump(tree))
def test_add_existing_import(self):
tree = ast.parse('from a.b import c')
self.assertEqual('c', import_utils.add_import(tree, 'a.b.c'))
self.assertEqual('from a.b import c\n', pasta.dump(tree))
def test_add_single_name_from_import(self):
tree = ast.parse('')
self.assertEqual(
'foo', import_utils.add_import(tree, 'foo', from_import=True))
self.assertEqual('import foo\n', pasta.dump(tree))
def test_add_from_import(self):
tree = ast.parse('')
self.assertEqual(
'c', import_utils.add_import(tree, 'a.b.c', from_import=True))
self.assertEqual('from a.b import c\n', pasta.dump(tree))
def test_add_normal_import(self):
tree = ast.parse('')
self.assertEqual(
'a.b.c', import_utils.add_import(tree, 'a.b.c', from_import=False))
self.assertEqual('import a.b.c\n', pasta.dump(tree))
