def test_add_import_preserve_doctring_multiples(self) -> None:
"""
Should preserve any doctring if adding to the beginning.
"""
before = """
# This is some docstring
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
# This is some docstring
import argparse
from a.b.c import D
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[ImportItem("a.b.c", "D", None), ImportItem("argparse", None, None)],
)
def test_add_explicit_relative_object_simple(self) -> None:
"""
Should add a relative import.
"""
before = """
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from ..a import B
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[ImportItem("a", "B", None, 2)],
)
def test_dont_add_explicit_relative_object_simple(self) -> None:
"""
Should not add object as an import since it exists.
"""
before = """
from ..c import D
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from ..c import D
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[ImportItem("c", "D", None, 2)],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def test_add_object_resolve_dotted_relative_modify_simple(self) -> None:
"""
Should merge a relative new module with an absolute existing one.
"""
before = """
from ..c import E, F
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from ..c import D, E, F
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[ImportItem("..c", "D", None)],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def test_add_object_explicit_relative_modify_simple(self) -> None:
"""
Should modify existing import to add new object.
"""
before = """
from ..c import E, F
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from ..c import D, E, F
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[ImportItem("c", "D", None, 2)],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def test_add_object_alias_modify_simple(self) -> None:
"""
Should modify existing import with alias to add new object
"""
before = """
from a.b.c import E, F
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from a.b.c import D as _, E, F
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(before, after, [ImportItem("a.b.c", "D", "_")])
def test_dont_add_object_alias_simple(self) -> None:
"""
Should not add object as an import since it exists.
"""
before = """
from a.b.c import D as E
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from a.b.c import D as E
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(before, after, [ImportItem("a.b.c", "D", "E")])
def test_add_future(self) -> None:
"""
Should add future import before any other imports.
"""
before = """
import unittest
import abc
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from __future__ import dummy_feature
import unittest
import abc
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before, after, [ImportItem("__future__", "dummy_feature", None)]
)
def test_add_explicit_relative_alias(self) -> None:
"""
Should add a relative import from .. .
"""
before = """
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from .. import a as foo
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[ImportItem("a", None, "foo", 2)],
)
def test_dont_add_module_simple(self) -> None:
"""
Should not add module as an import since it exists
"""
before = """
import a.b.c
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
import a.b.c
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(before, after, [ImportItem("a.b.c", None, None)])
def test_add_module_complex(self) -> None:
"""
Should add some modules as an import.
"""
before = """
import argparse
import sys
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
import argparse
import sys
import a.b.c
import defg.hi
import jkl as h
import i.j as k
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[
ImportItem("a.b.c", None, None),
ImportItem("defg.hi", None, None),
ImportItem("argparse", None, None),
ImportItem("jkl", None, "h"),
ImportItem("i.j", None, "k"),
],
)
def test_import_order(self) -> None:
"""
The imports should be in alphabetic order of added imports, added import alias, original imports.
"""
before = """
from a import b, e, h
"""
after = """
from a import c, f, d as x, g as y, b, e, h
"""
self.assertCodemod(
before,
after,
[
ImportItem("a", "f", None),
ImportItem("a", "g", "y"),
ImportItem("a", "c", None),
ImportItem("a", "d", "x"),
],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def test_strict_module_no_imports(self) -> None:
"""
First added import in strict module should go after __strict__ flag.
"""
before = """
__strict__ = True
class Foo:
pass
"""
after = """
__strict__ = True
import argparse
class Foo:
pass
"""
self.assertCodemod(before, after, [ImportItem("argparse", None, None)])
def test_import_in_docstring_module(self) -> None:
"""
The import should be added after module docstring.
"""
before = """
'''Docstring.'''
import typing
"""
after = """
'''Docstring.'''
from __future__ import annotations
import typing
"""
self.assertCodemod(
before,
after,
[ImportItem("__future__", "annotations", None)],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def test_add_module_alias_simple(self) -> None:
"""
Should add module with alias as an import.
"""
before = """
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
import a.b.c as d
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(before, after, [ImportItem("a.b.c", None, "d")])
def test_add_object_simple(self) -> None:
"""
Should add object as an import.
"""
before = """
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from a.b.c import D
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(before, after, [ImportItem("a.b.c", "D", None)])
def test_add_object_modify_complex(self) -> None:
"""
Should modify existing import to add new object
"""
before = """
from a.b.c import E, F, G as H
from d.e.f import Foo, Bar
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from a.b.c import D, E, F, G as H
from d.e.f import Baz as Qux, Foo, Bar
from g.h.i import V as W, X, Y, Z
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[
ImportItem("a.b.c", "D", None),
ImportItem("a.b.c", "F", None),
ImportItem("a.b.c", "G", "H"),
ImportItem("d.e.f", "Foo", None),
ImportItem("g.h.i", "Z", None),
ImportItem("g.h.i", "X", None),
ImportItem("d.e.f", "Bar", None),
ImportItem("d.e.f", "Baz", "Qux"),
ImportItem("g.h.i", "Y", None),
ImportItem("g.h.i", "V", "W"),
ImportItem("a.b.c", "F", None),
],
)
def test_add_and_modify_complex(self) -> None:
"""
Should correctly add both module and object imports
"""
before = """
import argparse
import sys
from a.b.c import E, F
from d.e.f import Foo, Bar
import bar as baz
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
import argparse
import sys
from a.b.c import D, E, F
from d.e.f import Foo, Bar
import bar as baz
import foo
import qux as quux
from g.h.i import X, Y, Z
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[
ImportItem("a.b.c", "D", None),
ImportItem("a.b.c", "F", None),
ImportItem("d.e.f", "Foo", None),
ImportItem("sys", None, None),
ImportItem("g.h.i", "Z", None),
ImportItem("g.h.i", "X", None),
ImportItem("d.e.f", "Bar", None),
ImportItem("g.h.i", "Y", None),
ImportItem("foo", None, None),
ImportItem("a.b.c", "F", None),
ImportItem("bar", None, "baz"),
ImportItem("qux", None, "quux"),
],
)