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_build_module_stubs_typed_dict_parameter(self):
function = FunctionDefinition.from_callable_and_traced_types(
Dummy.an_instance_method,
{
'foo': TypedDict(DUMMY_TYPED_DICT_NAME, {'a': int, 'b': str}),
'bar': int,
},
TypedDict(DUMMY_TYPED_DICT_NAME, {'c': int}),
None,
existing_annotation_strategy=ExistingAnnotationStrategy.IGNORE
)
entries = [function]
expected = module_stub_for_method_with_typed_dict
self.maxDiff = None
assert build_module_stubs(entries) == expected
def test_from_callable_and_traced_types(self, func, arg_types,
return_type, yield_type, expected):
function = FunctionDefinition.from_callable_and_traced_types(
func, arg_types,
return_type, yield_type,
existing_annotation_strategy=ExistingAnnotationStrategy.IGNORE
)
assert function == expected
def test_render_yield_typed_dict(self):
function = FunctionDefinition.from_callable_and_traced_types(
Dummy.an_instance_method,
{
'foo': int,
'bar': int,
},
int,
yield_type=TypedDict(DUMMY_TYPED_DICT_NAME, {'a': int, 'b': str}),
existing_annotation_strategy=ExistingAnnotationStrategy.IGNORE
)
entries = [function]
expected = '\n'.join([
'from mypy_extensions import TypedDict',
'from typing import Generator',
'',
'',
'class DummyAnInstanceMethodYieldTypedDict(TypedDict):',
def test_render_nested_typed_dict(self):
function = FunctionDefinition.from_callable_and_traced_types(
Dummy.an_instance_method,
{
'foo': TypedDict(DUMMY_TYPED_DICT_NAME,
{
# Naming the key 'z' to test a class name
# that comes last in alphabetical order.
'z': TypedDict(DUMMY_TYPED_DICT_NAME, {'a': int, 'b': str}),
'b': str,
}),
'bar': int,
},
int,
None,
existing_annotation_strategy=ExistingAnnotationStrategy.IGNORE
)
entries = [function]
def get_updated_definition(
func: Callable,
traces: Iterable[CallTrace],
rewriter: Optional[TypeRewriter] = None,
existing_annotation_strategy: ExistingAnnotationStrategy = ExistingAnnotationStrategy.REPLICATE,
) -> FunctionDefinition:
"""Update the definition for func using the types collected in traces."""
if rewriter is None:
rewriter = NoOpRewriter()
arg_types, return_type, yield_type = shrink_traced_types(traces)
arg_types = {name: rewriter.rewrite(typ) for name, typ in arg_types.items()}
if return_type is not None:
return_type = rewriter.rewrite(return_type)
if yield_type is not None:
yield_type = rewriter.rewrite(yield_type)
return FunctionDefinition.from_callable_and_traced_types(func, arg_types, return_type,
yield_type, existing_annotation_strategy)