def test_io_import_single(self):
"""Single _io imports should convert to io"""
imports = ImportMap()
imports['_io'] = {'BytesIO'}
stub = ImportBlockStub(imports)
expected = "\n".join([
'from io import BytesIO',
])
assert stub.render() == expected
def test_single_import(self):
"""Single imports should be on one line"""
imports = ImportMap()
imports['a.module'] = {'AClass'}
imports['another.module'] = {'AnotherClass'}
stub = ImportBlockStub(imports)
expected = "\n".join([
'from a.module import AClass',
'from another.module import AnotherClass',
])
assert stub.render() == expected
def test_multiple_io_imports(self):
"""Multiple imports from single _io module should be convert to io import"""
imports = ImportMap()
imports['_io'] = {'BytesIO', 'FileIO'}
stub = ImportBlockStub(imports)
expected = "\n".join([
'from io import (',
' BytesIO,',
' FileIO,',
')',
])
assert stub.render() == expected
def test_multiple_imports(self):
"""Multiple imports from a single module should each be on their own line"""
imports = ImportMap()
imports['a.module'] = {'AClass', 'AnotherClass', 'AThirdClass'}
stub = ImportBlockStub(imports)
expected = "\n".join([
'from a.module import (',
' AClass,',
' AThirdClass,',
' AnotherClass,',
')',
])
assert stub.render() == expected
def test_merge(self):
a = ImportMap()
a['module.a'] = {'ClassA', 'ClassB'}
a['module.b'] = {'ClassE', 'ClassF'}
b = ImportMap()
b['module.a'] = {'ClassB', 'ClassC'}
b['module.c'] = {'ClassX', 'ClassY'}
expected = ImportMap()
for mod in ('module.a', 'module.b', 'module.c'):
expected[mod] = a[mod] | b[mod]
a.merge(b)
assert a == expected
ClassStub(
name='FooBar2TypedDict(TypedDict)',
function_stubs=[],
attribute_stubs=[
AttributeStub(name='b', typ=str),
])]),
),
],
)
def test_replace_typed_dict_with_stubs(self, typ, expected):
rewritten_type, stubs = ReplaceTypedDictsWithStubs.rewrite_and_get_stubs(typ, class_name_hint='foo_bar')
actual = rewritten_type, stubs
assert actual == expected
typed_dict_import_map = ImportMap()
typed_dict_import_map['mypy_extensions'] = {'TypedDict'}
module_stub_for_method_with_typed_dict = {
'tests.util': ModuleStub(
function_stubs=(),
class_stubs=[
ClassStub(
name='Dummy',
function_stubs=[
FunctionStub(
name='an_instance_method',
signature=Signature(
parameters=[
Parameter(name='self',
kind=Parameter.POSITIONAL_OR_KEYWORD,
annotation=Parameter.empty),
