fix the rewriter on relative imports (fixes #58)

This commit is contained in:
Benjamin Peterson 2011-06-28 10:39:11 -05:00
parent f4eb15632d
commit c6e3606c6b
2 changed files with 9 additions and 2 deletions

View File

@ -109,8 +109,8 @@ class AssertionRewriter(ast.NodeVisitor):
return return
lineno += len(doc) - 1 lineno += len(doc) - 1
expect_docstring = False expect_docstring = False
elif (not isinstance(item, ast.ImportFrom) or item.level > 0 and elif (not isinstance(item, ast.ImportFrom) or item.level > 0 or
item.identifier != "__future__"): item.module != "__future__"):
lineno = item.lineno lineno = item.lineno
break break
pos += 1 pos += 1

View File

@ -76,6 +76,13 @@ class TestAssertionRewrite:
assert imp.lineno == 3 assert imp.lineno == 3
assert imp.col_offset == 0 assert imp.col_offset == 0
assert isinstance(m.body[4], ast.Expr) assert isinstance(m.body[4], ast.Expr)
s = """from . import relative\nother_stuff"""
m = rewrite(s)
for imp in m.body[0:2]:
assert isinstance(imp, ast.Import)
assert imp.lineno == 1
assert imp.col_offset == 0
assert isinstance(m.body[3], ast.Expr)
def test_dont_rewrite(self): def test_dont_rewrite(self):
s = """'PYTEST_DONT_REWRITE'\nassert 14""" s = """'PYTEST_DONT_REWRITE'\nassert 14"""