def _infer_augmented_assign(node, context, solver):
"""Infer the types for augmented assignments
Examples:
a += 5
b[2] &= x
c.x -= f(1, 2)
"""
target_type = expr.infer(node.target, context, solver)
value_type = expr.infer(node.value, context, solver)
result_type = expr.binary_operation_type(target_type, node.op, value_type,
node.lineno, solver)
_infer_assignment_target(node.target, context, result_type, solver)
return solver.z3_types.none
def _infer_with(node, context, solver):
"""Infer the types for a with block"""
for item in node.items:
if item.optional_vars:
item_type = expr.infer(item.context_expr, context, solver)
_infer_assignment_target(item.optional_vars, context, item_type,
solver)
return _infer_body(node.body, context, node.lineno, solver)
def infer(node, context, solver, parent=None):
if parent:
node._parent = parent
if isinstance(node, ast.Assign):
return _infer_assign(node, context, solver)
elif isinstance(node, ast.AugAssign):
return _infer_augmented_assign(node, context, solver)
elif isinstance(node, ast.Return):
if not node.value:
return solver.z3_types.none
return expr.infer(node.value, context, solver)
elif isinstance(node, ast.Delete):
return _infer_delete(node, context, solver)
elif isinstance(node, (ast.If, ast.While)):
return _infer_control_flow(node, context, solver)
elif isinstance(node, ast.For):
return _infer_for(node, context, solver)
elif sys.version_info[0] >= 3 and sys.version_info[1] >= 5 and isinstance(
node, ast.AsyncFor):
# AsyncFor is introduced in Python 3.5
return _infer_for(node, context, solver)
elif isinstance(node, ast.With):
return _infer_with(node, context, solver)
elif sys.version_info[0] >= 3 and sys.version_info[1] >= 5 and isinstance(
node, ast.AsyncWith):
# AsyncWith is introduced in Python 3.5
return _infer_with(node, context, solver)
elif isinstance(node, ast.Try):
return _infer_try(node, context, solver)
elif isinstance(node, ast.FunctionDef):
return _infer_func_def(node, context, solver)
elif isinstance(node, ast.ClassDef):
return _infer_class_def(node, context, solver)
elif isinstance(node, ast.Expr):
expr.infer(node.value, context, solver)
elif isinstance(node, ast.Import):
return _infer_import(node, context, solver)
elif isinstance(node, ast.ImportFrom):
return _infer_import_from(node, context, solver)
elif isinstance(node, ast.Raise):
return _infer_raise(node, context, solver)
elif sys.version_info >= (3, 6) and isinstance(node, ast.AnnAssign):
return _infer_annotated_assign(node, context, solver)
return solver.z3_types.none
def _delete_element(target, context, lineno, solver):
"""Remove (if needed) a target from the context
Cases:
- del var_name: remove its type mapping from the context directly.
- del subscript:
* Tuple/String --> Immutable. Raise exception.
* List/Dict --> Do nothing to the context.
"""
if isinstance(target, (ast.Tuple, ast.List)): # Multiple deletions
for elem in target.elts:
_delete_element(elem, context, lineno, solver)
elif isinstance(target, ast.Name):
context.delete_type(target.id)
elif isinstance(target, ast.Subscript):
expr.infer(target, context, solver)
indexed_type = expr.infer(target.value, context, solver)
solver.add(axioms.delete_subscript(indexed_type, solver.z3_types),
fail_message="Deletion in line {}".format(lineno))
elif isinstance(target, ast.Attribute):
raise NotImplementedError("Attribute deletion is not supported.")
def _infer_annotated_assign(node, context, solver):
"""Infer the types of variables in annotated assignment node"""
if node.value:
value_type = expr.infer(node.value, context, solver)
target_type = _infer_assignment_target(node.target, context,
value_type, solver)
else:
target_type = _infer_one_target(node.target, context, solver)
annotation_type = solver.resolve_annotation(node.annotation,
get_module(node))
solver.add(target_type == annotation_type,
fail_message="Annotated assignment in line {}".format(
node.lineno))
def _infer_one_target(target, context, solver):
"""
Get the type of the left hand side of an assignment
:param target: The target on the left hand side
:param context: The current context level
:param solver: The SMT solver
:return: the type of the target
"""
if isinstance(target, ast.Name):
if target.id in context.types_map:
return context.get_type(target.id)
else:
target_type = solver.new_z3_const("assign")
context.set_type(target.id, target_type)
return target_type
elif isinstance(target, ast.Tuple):
args_types = []
for elt in target.elts:
args_types.append(_infer_one_target(elt, context, solver))
return solver.z3_types.tuples[len(args_types)](*args_types)
elif isinstance(target, ast.List):
list_type = solver.new_z3_const("assign")
for elt in target.elts:
solver.add(list_type == _infer_one_target(elt, context, solver),
fail_message="List assignment in line {}".format(
target.lineno))
return solver.z3_types.list(list_type)
target_type = expr.infer(target, context, solver)
if isinstance(target, ast.Subscript):
solver.add(axioms.subscript_assignment(
expr.infer(target.value, context, solver), solver.z3_types),
fail_message="Subscript assignment in line {}".format(
target.lineno))
return target_type
def _infer_assign(node, context, solver):
"""Infer the types of target variables in an assignment node."""
if _is_type_var_declaration(node.value):
module = get_module(node)
solver.annotation_resolver.add_type_var(node.targets[0], node.value,
solver, module)
else:
value_type = expr.infer(node.value, context, solver)
for target in node.targets:
target_type = _infer_assignment_target(target, context, value_type,
solver)
if len(node.targets) == 1:
context.add_assignment(target_type, node)
return solver.z3_types.none
def _infer_args_defaults(args_types, defaults, context, solver):
"""Infer the default values of function arguments (if any)
:param args_types: Z3 constants for arguments types
:param defaults: AST nodes for default values of arguments
:param context: The parent of the function context
:param solver: The inference Z3 solver
A default array of length `n` represents the default values of the last `n` arguments
"""
for i, default in enumerate(defaults):
arg_idx = i + len(args_types) - len(
defaults) # The defaults array correspond to the last arguments
default_type = expr.infer(default, context, solver)
solver.add(solver.z3_types.subtype(default_type, args_types[arg_idx]),
fail_message="Function default argument in line {}".format(
defaults[i].lineno))
solver.optimize.add_soft(default_type == args_types[arg_idx])
def _infer_for(node, context, solver):
"""Infer the type for a for loop node
Limitation:
- The iterable can't be a tuple.
For example: the following is not allowed:
for x in (1, 2.0, "string"):
....
"""
iter_type = expr.infer(node.iter, context, solver)
# Infer the target in the loop, inside the global context
# Cases:
# - Var name. Ex: for i in range(5)..
# - Tuple. Ex: for (a,b) in [(1,"st"), (3,"st2")]..
# - List. Ex: for [a,b] in [(1, "st"), (3, "st2")]..
target_type = solver.new_z3_const("for_target")
solver.add(axioms.for_loop(iter_type, target_type, solver.z3_types),
fail_message="For loop in line {}".format(node.lineno))
_infer_assignment_target(node.target, context, target_type, solver)
return _infer_control_flow(node, context, solver)
def _infer_raise(node, context, solver):
if isinstance(node.exc, ast.Call):
for arg in node.exc.args:
expr.infer(arg, context, solver)
return solver.z3_types.none
def _infer_control_flow(node, context, solver):
"""Infer the type(s) for an if/while/for statements block.
Arguments:
node: The AST node to be inferred
context: the current context level
Example:
if (some_condition):
......
return "some string"
else:
......
return 2.0
type: super{String, Float}
"""
body_context = Context(node, node.body, solver, parent_context=context)
else_context = Context(node, node.body, solver, parent_context=context)
var_is_instance = None
if hasattr(node, "test"):
expr.infer(node.test, context, solver)
if (isinstance(node.test, ast.Call)
and isinstance(node.test.func, ast.Name)
and node.test.func.id == "isinstance"
and len(node.test.args) == 2):
# isinstance(x, t) test: Make `x` to be of type `t` in the then branch
# The only allowed case is when `x` is variable and `t` is a single type (not a tuple).
# Get the type `t`
t = solver.resolve_annotation(node.test.args[1], get_module(node))
# Set `x` to be an instance of `t` in the then branch
body_context.isinstance_nodes[ast.dump(node.test.args[0],
annotate_fields=False)] = t
# Keep track of the name of the variable `x`
if isinstance(node.test.args[0], ast.Name):
var_is_instance = node.test.args[0].id
body_type = _infer_body(node.body, body_context, node.lineno, solver)
else_type = _infer_body(node.orelse, else_context, node.lineno, solver)
# Re-assigning variables in the body branch
for v in body_context.types_map:
if context.has_variable(v) and v != var_is_instance:
t1 = body_context.types_map[v]
t2 = context.get_type(v)
solver.add(
t1 == t2,
fail_message="re-assigning in flow branching in line {}".
format(node.lineno))
# Re-assigning variables in the else branch
for v in else_context.types_map:
if context.has_variable(v):
t1 = else_context.types_map[v]
t2 = context.get_type(v)
solver.add(
t1 == t2,
fail_message="re-assigning in flow branching in line {}".
format(node.lineno))
# Take intersection of variables in both contexts
for v in body_context.types_map:
if v in else_context.types_map and not context.has_variable(v):
var_type = solver.new_z3_const("branching_var")
t1 = body_context.types_map[v]
t2 = else_context.types_map[v]
if inference_config["enforce_same_type_in_branches"]:
branch_axioms = [t1 == var_type, t2 == var_type]
else:
branch_axioms = [
solver.z3_types.subtype(t1, var_type),
solver.z3_types.subtype(t2, var_type)
]
solver.add(
branch_axioms,
fail_message="subtyping in flow branching in line {}".format(
node.lineno))
solver.optimize.add_soft(t1 == var_type)
solver.optimize.add_soft(t2 == var_type)
context.set_type(v, var_type)
result_type = solver.new_z3_const("control_flow")
solver.add(axioms.control_flow(body_type, else_type, result_type,
solver.z3_types),
fail_message="Control flow in line {}".format(node.lineno))
solver.optimize.add_soft(result_type == body_type)
solver.optimize.add_soft(result_type == else_type)
return result_type