Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_render(self):
cm_stub = _func_stub_from_callable(Dummy.a_class_method)
im_stub = _func_stub_from_callable(Dummy.an_instance_method)
func_stubs = (cm_stub, im_stub)
test_stub = ClassStub('Test', function_stubs=func_stubs)
test2_stub = ClassStub('Test2', function_stubs=func_stubs)
other_class_stubs = module_stub_for_method_with_typed_dict['tests.util'].class_stubs.values()
class_stubs = (*other_class_stubs, test_stub, test2_stub)
typed_dict_class_stubs = module_stub_for_method_with_typed_dict['tests.util'].typed_dict_class_stubs
mod_stub = ModuleStub(function_stubs=func_stubs,
class_stubs=class_stubs,
typed_dict_class_stubs=typed_dict_class_stubs)
expected = '\n'.join([
'class DummyAnInstanceMethodTypedDict(TypedDict):',
' c: int',
'',
'',
'class FooTypedDict(TypedDict):',
' a: int',
' b: str',
'',
'',
def test_render(self):
cm_stub = _func_stub_from_callable(Dummy.a_class_method.__func__)
im_stub = _func_stub_from_callable(Dummy.an_instance_method)
class_stub = ClassStub('Test', function_stubs=(cm_stub, im_stub),
attribute_stubs=[
AttributeStub('foo', int),
AttributeStub('bar', str),
])
expected = '\n'.join([
'class Test:',
' bar: str',
' foo: int',
' @classmethod',
' def a_class_method(cls, foo: Any) -> Optional[frame]: ...',
' def an_instance_method(self, foo: Any, bar: Any) -> Optional[frame]: ...',
])
assert class_stub.render() == expected
AttributeStub(name='b', typ=str),
])],
),
],
)
def test_stubs_from_typed_dict(self, typ, expected):
assert ClassStub.stubs_from_typed_dict(typ, class_name='FooBar') == 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),
Parameter(name='foo',
kind=Parameter.POSITIONAL_OR_KEYWORD,
annotation=make_forward_ref('FooTypedDict')),
Parameter(name='bar',
kind=Parameter.POSITIONAL_OR_KEYWORD,
annotation=int),
],
def test_render(self):
cm_stub = _func_stub_from_callable(Dummy.a_class_method)
im_stub = _func_stub_from_callable(Dummy.an_instance_method)
func_stubs = (cm_stub, im_stub)
test_stub = ClassStub('Test', function_stubs=func_stubs)
test2_stub = ClassStub('Test2', function_stubs=func_stubs)
other_class_stubs = module_stub_for_method_with_typed_dict['tests.util'].class_stubs.values()
class_stubs = (*other_class_stubs, test_stub, test2_stub)
typed_dict_class_stubs = module_stub_for_method_with_typed_dict['tests.util'].typed_dict_class_stubs
mod_stub = ModuleStub(function_stubs=func_stubs,
class_stubs=class_stubs,
typed_dict_class_stubs=typed_dict_class_stubs)
expected = '\n'.join([
'class DummyAnInstanceMethodTypedDict(TypedDict):',
' c: int',
'',
'',
'class FooTypedDict(TypedDict):',
' a: int',
' b: str',
'',
if len(class_path) > 0:
klass = '.'.join(class_path)
if entry.module not in mod_stubs:
mod_stubs[entry.module] = ModuleStub()
mod_stub = mod_stubs[entry.module]
imports = get_imports_for_signature(entry.signature)
# Import TypedDict, if needed.
if entry.typed_dict_class_stubs:
imports['mypy_extensions'].add('TypedDict')
func_stub = FunctionStub(name, entry.signature, entry.kind, list(imports.keys()), entry.is_async)
# Don't need to import anything from the same module
imports.pop(entry.module, None)
mod_stub.imports_stub.imports.merge(imports)
if klass is not None:
if klass not in mod_stub.class_stubs:
mod_stub.class_stubs[klass] = ClassStub(klass)
class_stub = mod_stub.class_stubs[klass]
class_stub.function_stubs[func_stub.name] = func_stub
else:
mod_stub.function_stubs[func_stub.name] = func_stub
mod_stub.typed_dict_class_stubs.extend(entry.typed_dict_class_stubs)
return mod_stubs
def __init__(
self,
function_stubs: Iterable[FunctionStub] = None,
class_stubs: Iterable[ClassStub] = None,
imports_stub: ImportBlockStub = None,
typed_dict_class_stubs: Iterable[ClassStub] = None,
) -> None:
self.function_stubs: Dict[str, FunctionStub] = {}
if function_stubs is not None:
self.function_stubs = {stub.name: stub for stub in function_stubs}
self.class_stubs: Dict[str, ClassStub] = {}
if class_stubs is not None:
self.class_stubs = {stub.name: stub for stub in class_stubs}
self.imports_stub = imports_stub if imports_stub else ImportBlockStub()
self.typed_dict_class_stubs: List[ClassStub] = []
if typed_dict_class_stubs is not None:
self.typed_dict_class_stubs = list(typed_dict_class_stubs)
if is_anonymous_typed_dict(typ):
class_name = get_typed_dict_class_name(name)
typed_dict_class_stubs.extend(ClassStub.stubs_from_typed_dict(typ, class_name))
typ = make_forward_ref(class_name)
new_arg_types[name] = typ
if return_type and is_anonymous_typed_dict(return_type):
# Replace the dot in a qualified name.
class_name = get_typed_dict_class_name(func.__qualname__.replace('.', '_'))
typed_dict_class_stubs.extend(ClassStub.stubs_from_typed_dict(return_type, class_name))
return_type = make_forward_ref(class_name)
if yield_type and is_anonymous_typed_dict(yield_type):
# Replace the dot in a qualified name.
class_name = get_typed_dict_class_name(func.__qualname__.replace('.', '_') + 'Yield')
typed_dict_class_stubs.extend(ClassStub.stubs_from_typed_dict(yield_type, class_name))
yield_type = make_forward_ref(class_name)
function = FunctionDefinition.from_callable(func)
signature = function.signature
signature = update_signature_args(signature, new_arg_types,
function.has_self, existing_annotation_strategy)
signature = update_signature_return(signature, return_type,
yield_type, existing_annotation_strategy)
return FunctionDefinition(function.module, function.qualname,
function.kind, signature,
function.is_async, typed_dict_class_stubs)
def stubs_from_typed_dict(type_dict: type, class_name: str) -> List['ClassStub']:
"""Return a list of class stubs for all TypedDicts found within `type_dict`."""
assert is_anonymous_typed_dict(type_dict)
class_stubs = []
attribute_stubs = []
for name, typ in type_dict.__annotations__.items():
if is_anonymous_typed_dict(typ):
_class_name = get_typed_dict_class_name(name)
class_stubs.extend(ClassStub.stubs_from_typed_dict(typ, _class_name))
typ = make_forward_ref(_class_name)
attribute_stubs.append(AttributeStub(name, typ))
class_stubs.append(ClassStub(name=f'{class_name}(TypedDict)',
function_stubs=[],
attribute_stubs=attribute_stubs))
return class_stubs