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_build_module_stubs(self):
entries = [
FunctionDefinition.from_callable(Dummy.a_static_method),
FunctionDefinition.from_callable(Dummy.a_class_method.__func__),
FunctionDefinition.from_callable(Dummy.an_instance_method),
FunctionDefinition.from_callable(simple_add),
]
simple_add_stub = _func_stub_from_callable(simple_add)
to_strip = ['typing']
dummy_stub = ClassStub(
'Dummy',
function_stubs=[
_func_stub_from_callable(Dummy.a_class_method.__func__,
to_strip),
_func_stub_from_callable(Dummy.an_instance_method, to_strip),
_func_stub_from_callable(Dummy.a_static_method, to_strip),
])
imports = {'typing': {'Any', 'Optional'}}
expected = {
'tests.test_stubs':
ModuleStub(function_stubs=[simple_add_stub]),
'tests.util':
ModuleStub(class_stubs=[dummy_stub],
imports_stub=ImportBlockStub(imports)),
}
self.maxDiff = None
assert build_module_stubs(entries) == expected
kind=Parameter.POSITIONAL_OR_KEYWORD,
annotation=make_forward_ref('FooTypedDict')),
Parameter(name='bar',
kind=Parameter.POSITIONAL_OR_KEYWORD,
annotation=int),
],
return_annotation=make_forward_ref('DummyAnInstanceMethodTypedDict'),
),
kind=FunctionKind.INSTANCE,
strip_modules=['mypy_extensions'],
is_async=False,
),
],
),
],
imports_stub=ImportBlockStub(typed_dict_import_map),
typed_dict_class_stubs=[
ClassStub(
name='FooTypedDict(TypedDict)',
function_stubs=[],
attribute_stubs=[
AttributeStub('a', int),
AttributeStub('b', str),
]
),
ClassStub(
# We use the name of the method, `Dummy.an_instance_method`,
# to get `DummyAnInstanceMethodTypedDict`.
name='DummyAnInstanceMethodTypedDict(TypedDict)',
function_stubs=[],
attribute_stubs=[