625 lines
19 KiB
Python
625 lines
19 KiB
Python
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Type inference.
|
|
|
|
This analysis annotates all symbols nodes of an AST with type information
|
|
extracted from static sources:
|
|
* type annotations
|
|
* global and local symbols visible to the function at analysis time
|
|
* literals
|
|
|
|
Important: This analysis is static, and does not detect dynamic type changes.
|
|
The analysis attempts to use the values of external symbols, if available. These
|
|
values are also considered static for the purpose of analysis.
|
|
|
|
Requires reaching function definitions analysis.
|
|
"""
|
|
|
|
import itertools
|
|
|
|
from typing import Any, Callable, Dict, Set
|
|
|
|
import gast
|
|
|
|
from tensorflow.python.autograph.pyct import anno
|
|
from tensorflow.python.autograph.pyct import cfg
|
|
from tensorflow.python.autograph.pyct import qual_names
|
|
from tensorflow.python.autograph.pyct import transformer
|
|
from tensorflow.python.autograph.pyct.static_analysis import activity
|
|
from tensorflow.python.autograph.pyct.static_analysis import annos
|
|
|
|
|
|
class Resolver(object):
|
|
"""Resolver objects handle the process of looking up actual names and types.
|
|
|
|
Unless noted otherwise, all resolve_* methods:
|
|
* have a first namespace argument, mapping string to actual values
|
|
* have a second types_namespace argument, mapping string to actual inferred
|
|
types
|
|
* specify names as QN objects
|
|
* specify types as a Set of inferred types
|
|
|
|
Unless noted otherwise, all resolve_* methods must return either:
|
|
* a set of `type` objects
|
|
* None
|
|
"""
|
|
|
|
def res_name(self, ns, types_ns, name):
|
|
"""Resolves the type/value an external (e.g. closure, global) variable.
|
|
|
|
Args:
|
|
ns: namespace
|
|
types_ns: types namespace
|
|
name: symbol name
|
|
Returns:
|
|
Tuple (type, static_value). The first element is the type to use for
|
|
inferrence. The second is the static value to use. Return None to treat it
|
|
as unknown.
|
|
"""
|
|
raise NotImplementedError('subclasses must implement')
|
|
|
|
def res_value(self, ns, value):
|
|
"""Resolves the type a literal or static value."""
|
|
raise NotImplementedError('subclasses must implement')
|
|
|
|
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
|
|
"""Resolves the type of a (possibly annotated) function argument.
|
|
|
|
Args:
|
|
ns: namespace
|
|
types_ns: types namespace
|
|
f_name: str, the function name
|
|
name: str, the argument name
|
|
type_anno: the type annotating the argument, if any
|
|
f_is_local: bool, whether the function is a local function
|
|
Returns:
|
|
Set of the argument types.
|
|
"""
|
|
raise NotImplementedError('subclasses must implement')
|
|
|
|
def res_call(self, ns, types_ns, node, f_type, args, keywords):
|
|
"""Resolves the return type an external function or method call.
|
|
|
|
Args:
|
|
ns: namespace
|
|
types_ns: types namespace
|
|
node: str, the function name
|
|
f_type: types of the actual function being called, if known
|
|
args: types of each respective argument in node.args
|
|
keywords: types of each respective argument in node.keywords
|
|
|
|
Returns:
|
|
Tuple (return_type, side_effect_types). The first element is just the
|
|
return types of the function. The second element is a map from
|
|
argument names to sets of types, and allow modelling side effects of
|
|
functions (for example via global or nonlocal).
|
|
"""
|
|
raise NotImplementedError('subclasses must implement')
|
|
|
|
# TODO(mdan): Clean this up.
|
|
def res_slice(self, ns, types_ns, node_or_slice, value, slice_):
|
|
"""Resolves the return type of slice operation."""
|
|
raise NotImplementedError('subclasses must implement')
|
|
|
|
def res_compare(self, ns, types_ns, node, left, right):
|
|
"""Resolves the return type of a unary operation."""
|
|
raise NotImplementedError('subclasses must implement')
|
|
|
|
def res_unop(self, ns, types_ns, node, opnd):
|
|
"""Resolves the return type of a unary operation."""
|
|
raise NotImplementedError('subclasses must implement')
|
|
|
|
def res_binop(self, ns, types_ns, node, left, right):
|
|
"""Resolves the return type of a binary operation."""
|
|
raise NotImplementedError('subclasses must implement')
|
|
|
|
def res_list_literal(self, ns, elt_types):
|
|
"""Resolves the type of a list literal from its elements."""
|
|
raise NotImplementedError('subclasses must implement')
|
|
|
|
|
|
class _TypeMap(object):
|
|
"""Abstraction for the state of the CFG walk for type inference.
|
|
|
|
This is a value type. Only implements the strictly necessary operators.
|
|
|
|
Attributes:
|
|
types: Dict[qual_names.QN, Set[Type]], mapping symbols to the set of
|
|
possible types.
|
|
"""
|
|
|
|
def __init__(self, init_from=None):
|
|
if init_from:
|
|
assert isinstance(init_from, _TypeMap)
|
|
self.types = {
|
|
s: set(other_types) for s, other_types in init_from.types.items()
|
|
}
|
|
else:
|
|
self.types = {}
|
|
|
|
def __eq__(self, other):
|
|
if frozenset(self.types.keys()) != frozenset(other.types.keys()):
|
|
return False
|
|
ret = all(self.types[s] == other.types[s] for s in self.types)
|
|
return ret
|
|
|
|
def __ne__(self, other):
|
|
return not self.__eq__(other)
|
|
|
|
def __or__(self, other):
|
|
assert isinstance(other, _TypeMap)
|
|
result = _TypeMap(self)
|
|
for s, other_types in other.types.items():
|
|
if s not in result.types:
|
|
self_types = set()
|
|
result.types[s] = self_types
|
|
else:
|
|
self_types = result.types[s]
|
|
self_types.update(other_types)
|
|
return result
|
|
|
|
def __repr__(self):
|
|
return 'SymbolTable {}'.format(self.types)
|
|
|
|
|
|
NO_VALUE = object()
|
|
|
|
|
|
class StmtInferrer(gast.NodeVisitor):
|
|
"""Runs type inference on a single AST statement.
|
|
|
|
This visitor annotates most nodes with type information. It also sets types
|
|
for the symbols modified by this statement in its types_out property.
|
|
|
|
Note: this inferrer is able to capture side effects of functions, however,
|
|
these side effects will not be applied to the current expression. Doing so
|
|
would create too much of a dependence on the runtime's internal rules about
|
|
execution order.
|
|
Example:
|
|
|
|
def f():
|
|
nonlocal a
|
|
a = 1
|
|
return a
|
|
|
|
a = 0.0
|
|
b = f() + a # a = float; side effect of f() ignored
|
|
print(a) # a = int; side effect of f() accounted for
|
|
"""
|
|
|
|
def __init__(self,
|
|
resolver: Resolver,
|
|
scope: activity.Scope,
|
|
namespace: Dict[qual_names.QN, Any],
|
|
closure_types: Dict[qual_names.QN, Set[Any]],
|
|
types_in: _TypeMap):
|
|
self.resolver = resolver
|
|
self.scope = scope
|
|
self.namespace = namespace
|
|
self.closure_types = closure_types
|
|
self.types_in = types_in
|
|
self.new_symbols = {}
|
|
|
|
# rvalue type. This property is set when encountering an assign operation,
|
|
# so that visiting nodes with Store ctx (typically found on left side of
|
|
# assignments) can infer the type they should receive.
|
|
self.rtype = None
|
|
|
|
def visit(self, node):
|
|
types = super().visit(node)
|
|
if __debug__:
|
|
self._check_set(types)
|
|
if types is not None:
|
|
# TODO(mdan): Normalize by removing subtypes.
|
|
anno.setanno(node, anno.Static.TYPES, tuple(types))
|
|
return types
|
|
|
|
def _check_set(self, value):
|
|
if value is not None and not isinstance(value, set):
|
|
raise ValueError('{} method expected to return set, got {}'.format(
|
|
self.resolver, value))
|
|
|
|
def visit_Constant(self, node):
|
|
types = self.resolver.res_value(self.namespace, node.value)
|
|
if __debug__:
|
|
self._check_set(types)
|
|
return types
|
|
|
|
def _apply_unpacking(self, node):
|
|
assert isinstance(node.ctx, gast.Store)
|
|
if self.rtype is not None:
|
|
original_stype = self.rtype
|
|
# TODO(mdan): Find a better way to express unpacking.
|
|
i_type = self.resolver.res_value(self.namespace, 0)
|
|
for i, elt in enumerate(node.elts):
|
|
self.rtype = self.resolver.res_slice(
|
|
self.namespace, self.types_in.types, i, original_stype, i_type)
|
|
self.visit(elt)
|
|
self.rtype = original_stype
|
|
return original_stype
|
|
return None
|
|
|
|
def visit_Tuple(self, node):
|
|
if isinstance(node.ctx, gast.Load):
|
|
elt_types = ()
|
|
for elt in node.elts:
|
|
types_ = self.visit(elt)
|
|
if types_ is None:
|
|
return None
|
|
elt_types += (types_,)
|
|
return set(itertools.product(*elt_types))
|
|
return self._apply_unpacking(node)
|
|
|
|
def visit_List(self, node):
|
|
if isinstance(node.ctx, gast.Load):
|
|
elt_types = tuple(self.visit(elt) for elt in node.elts)
|
|
return self.resolver.res_list_literal(self.namespace, elt_types)
|
|
return self._apply_unpacking(node)
|
|
|
|
def visit_Set(self, node):
|
|
raise NotImplementedError()
|
|
|
|
def visit_Name(self, node):
|
|
name = anno.getanno(node, anno.Basic.QN)
|
|
|
|
if isinstance(node.ctx, gast.Load):
|
|
types = self.types_in.types.get(name, None)
|
|
if types is None:
|
|
if (name not in self.scope.bound) or (name in self.scope.nonlocals):
|
|
# TODO(mdan): Test with global variables.
|
|
if name in self.closure_types:
|
|
types = self.closure_types[name]
|
|
else:
|
|
types, value = self.resolver.res_name(
|
|
self.namespace, self.types_in.types, name)
|
|
if value is not None:
|
|
anno.setanno(node, anno.Static.VALUE, value)
|
|
|
|
elif isinstance(node.ctx, gast.Param):
|
|
# The direct parent it the whole function scope. See activity.py.
|
|
f_is_local = self.scope.parent.parent is not None
|
|
|
|
type_name = anno.getanno(node.annotation, anno.Basic.QN, None)
|
|
types = self.resolver.res_arg(self.namespace, self.types_in.types,
|
|
self.scope.function_name, name, type_name,
|
|
f_is_local)
|
|
if types is not None:
|
|
self.new_symbols[name] = types
|
|
|
|
elif isinstance(node.ctx, gast.Store):
|
|
if self.rtype is not None:
|
|
self.new_symbols[name] = self.rtype
|
|
types = self.rtype
|
|
|
|
else:
|
|
assert False, 'unknown ctx'
|
|
|
|
if __debug__:
|
|
self._check_set(types)
|
|
|
|
return types
|
|
|
|
def visit_Attribute(self, node):
|
|
parent_types = self.visit(node.value)
|
|
|
|
# Attempt to use the static value if known.
|
|
parent_value = anno.Static.VALUE.of(node.value, None)
|
|
if parent_value is not None:
|
|
static_value = getattr(parent_value, node.attr, NO_VALUE)
|
|
|
|
if static_value is NO_VALUE:
|
|
# Unexpected failure to resolve attribute. Ask the resolver about the
|
|
# full name instead.
|
|
types, static_value = self.resolver.res_name(
|
|
self.namespace, self.types_in, anno.Basic.QN.of(node))
|
|
anno.setanno(node, anno.Static.VALUE, static_value)
|
|
if __debug__:
|
|
self._check_set(types)
|
|
return types
|
|
|
|
else:
|
|
# Fall back to the type if that is known.
|
|
if parent_types is None:
|
|
return None
|
|
|
|
inferred_values = [getattr(t, node.attr, None) for t in parent_types]
|
|
if not inferred_values:
|
|
return None
|
|
|
|
static_value = inferred_values[0]
|
|
if static_value is None:
|
|
return None
|
|
|
|
if any(v is not static_value for v in inferred_values[1:]):
|
|
# Static value not stable, assume it's dynamic.
|
|
return None
|
|
|
|
types = self.resolver.res_value(self.namespace, static_value)
|
|
anno.setanno(node, anno.Static.VALUE, static_value)
|
|
|
|
if __debug__:
|
|
self._check_set(types)
|
|
|
|
return types
|
|
|
|
def visit_FunctionDef(self, node):
|
|
f_name = qual_names.QN(node.name)
|
|
|
|
if node.decorator_list:
|
|
raise NotImplementedError('decorators: {}'.format(node.decorator_list))
|
|
|
|
ret_types = None
|
|
if node.returns:
|
|
ret_types, _ = self.resolver.res_name(
|
|
self.namespace, self.types_in.types, anno.Basic.QN.of(node.returns))
|
|
if __debug__:
|
|
self._check_set(ret_types)
|
|
|
|
if ret_types is None:
|
|
ret_types = {Any}
|
|
|
|
f_types = set()
|
|
for rt in ret_types:
|
|
f_types.add(Callable[[Any], rt])
|
|
|
|
self.new_symbols[f_name] = f_types
|
|
# The definition of a function is an expression, hence has no return value.
|
|
return None
|
|
|
|
def _resolve_typed_callable(self, f_types, arg_types, keyword_types):
|
|
ret_types = set()
|
|
for t in f_types:
|
|
|
|
if isinstance(t, Callable):
|
|
# Note: these are undocummented - may be version-specific!
|
|
# Callable[[x], y]: __args__ are (x, y)
|
|
args = t.__args__
|
|
if args:
|
|
ret_types.add(args[-1])
|
|
else:
|
|
ret_types.add(Any)
|
|
else:
|
|
raise NotImplementedError('callable type {}'.format(type(t)))
|
|
|
|
# Side effects can not be inferred based on type alone.
|
|
side_effects = None
|
|
return ret_types, side_effects
|
|
|
|
def visit_Call(self, node):
|
|
self.visit(node.func)
|
|
|
|
f_name = anno.Basic.QN.of(node.func)
|
|
arg_types = [self.visit(a) for a in node.args]
|
|
keyword_types = [self.visit(kw.value) for kw in node.keywords]
|
|
|
|
if f_name in self.scope.bound:
|
|
# Local function, use local type definitions, if available.
|
|
f_type = self.types_in.types.get(f_name, None)
|
|
if f_type is None:
|
|
# No static type info available, nothing more to do.
|
|
ret_type, side_effects = None, None
|
|
else:
|
|
ret_type, side_effects = self._resolve_typed_callable(
|
|
f_type, arg_types, keyword_types)
|
|
|
|
else:
|
|
# Nonlocal function, resolve externally.
|
|
f_type = anno.Static.TYPES.of(node.func, None)
|
|
ret_type, side_effects = self.resolver.res_call(self.namespace,
|
|
self.types_in.types, node,
|
|
f_type, arg_types,
|
|
keyword_types)
|
|
|
|
if __debug__:
|
|
self._check_set(ret_type)
|
|
if side_effects:
|
|
if not isinstance(side_effects, dict):
|
|
raise ValueError(
|
|
'side effects must be dict, got {}'.format(side_effects))
|
|
for k, v in side_effects.items():
|
|
if not isinstance(k, qual_names.QN):
|
|
raise ValueError('side effect keys must be QNs, got {}'.format(k))
|
|
self._check_set(v)
|
|
|
|
if side_effects:
|
|
self.new_symbols.update(side_effects)
|
|
return ret_type
|
|
|
|
def visit_Expr(self, node):
|
|
return self.visit(node.value)
|
|
|
|
def visit_Assign(self, node):
|
|
self.rtype = self.visit(node.value)
|
|
|
|
for t in node.targets:
|
|
self.visit(t)
|
|
|
|
self.rtype = None
|
|
|
|
def visit_Subscript(self, node):
|
|
val_types = self.visit(node.value)
|
|
slice_types = self.visit(node.slice)
|
|
|
|
if val_types is None or slice_types is None:
|
|
return None
|
|
|
|
types = self.resolver.res_slice(
|
|
self.namespace, self.types_in.types, node, val_types, slice_types)
|
|
|
|
if __debug__:
|
|
self._check_set(types)
|
|
|
|
return types
|
|
|
|
def visit_Compare(self, node):
|
|
left_types = self.visit(node.left)
|
|
right_types = [self.visit(c) for c in node.comparators]
|
|
|
|
if left_types is None or any(t is None for t in right_types):
|
|
return None
|
|
|
|
types = self.resolver.res_compare(
|
|
self.namespace, self.types_in.types, node, left_types, right_types)
|
|
|
|
if __debug__:
|
|
self._check_set(types)
|
|
|
|
return types
|
|
|
|
def visit_BinOp(self, node):
|
|
left_types = self.visit(node.left)
|
|
right_types = self.visit(node.right)
|
|
|
|
if left_types is None or right_types is None:
|
|
return None
|
|
|
|
types = self.resolver.res_binop(
|
|
self.namespace, self.types_in.types, node, left_types, right_types)
|
|
|
|
if __debug__:
|
|
self._check_set(types)
|
|
|
|
return types
|
|
|
|
def visit_UnaryOp(self, node):
|
|
opnd_types = self.visit(node.operand)
|
|
|
|
if opnd_types is None:
|
|
return None
|
|
|
|
types = self.resolver.res_unop(
|
|
self.namespace, self.types_in.types, node, opnd_types)
|
|
|
|
if __debug__:
|
|
self._check_set(types)
|
|
|
|
return types
|
|
|
|
|
|
class Analyzer(cfg.GraphVisitor):
|
|
"""CFG visitor that propagates type information across statements."""
|
|
|
|
def __init__(self, graph, resolver, namespace, scope, closure_types):
|
|
"""Creates a new analyzer.
|
|
|
|
Args:
|
|
graph: cfg.Graph
|
|
resolver: Resolver
|
|
namespace: Dict[str, Any]
|
|
scope: activity.Scope
|
|
closure_types: Dict[QN, Set]
|
|
"""
|
|
super(Analyzer, self).__init__(graph)
|
|
self.resolver = resolver
|
|
self.namespace = namespace
|
|
self.scope = scope
|
|
self.closure_types = closure_types
|
|
|
|
context_types = {
|
|
n: t for n, t in closure_types.items() if n not in scope.bound
|
|
}
|
|
if context_types:
|
|
self.context_types = _TypeMap()
|
|
self.context_types.types = context_types
|
|
else:
|
|
self.context_types = None
|
|
|
|
def init_state(self, _):
|
|
return _TypeMap()
|
|
|
|
def _update_closure_types(self, ast_node, types):
|
|
existing_types = anno.Static.CLOSURE_TYPES.of(ast_node, None)
|
|
|
|
if existing_types is None:
|
|
existing_types = {}
|
|
anno.Static.CLOSURE_TYPES.add_to(ast_node, existing_types)
|
|
|
|
for k, v in types.types.items():
|
|
if k in existing_types:
|
|
existing_types[k].update(v)
|
|
else:
|
|
existing_types[k] = set(v)
|
|
|
|
def visit_node(self, node):
|
|
prev_types_out = self.out[node]
|
|
|
|
types_in = _TypeMap()
|
|
for n in node.prev:
|
|
types_in |= self.out[n]
|
|
if (self.context_types is not None) and (node is self.graph.entry):
|
|
types_in |= self.context_types
|
|
|
|
types_out = _TypeMap(types_in)
|
|
ast_node = node.ast_node
|
|
|
|
inferrer = StmtInferrer(self.resolver, self.scope, self.namespace,
|
|
self.closure_types, types_in)
|
|
inferrer.visit(ast_node)
|
|
types_out.types.update(inferrer.new_symbols)
|
|
|
|
reaching_fndefs = anno.Static.DEFINED_FNS_IN.of(ast_node)
|
|
node_scope = anno.Static.SCOPE.of(ast_node, None)
|
|
if node_scope is not None:
|
|
# TODO(mdan): Check that it's actually safe to skip nodes without scope.
|
|
reads = {str(qn) for qn in node_scope.read}
|
|
for def_node in reaching_fndefs:
|
|
if def_node.name in reads:
|
|
self._update_closure_types(def_node, types_out)
|
|
|
|
self.in_[node] = types_in
|
|
self.out[node] = types_out
|
|
|
|
return prev_types_out != types_out
|
|
|
|
|
|
class FunctionVisitor(transformer.Base):
|
|
"""AST visitor that applies type inference to each function separately."""
|
|
|
|
def __init__(self, source_info, graphs, resolver):
|
|
super(FunctionVisitor, self).__init__(source_info)
|
|
self.graphs = graphs
|
|
self.resolver = resolver
|
|
|
|
def visit_FunctionDef(self, node):
|
|
subgraph = self.graphs[node]
|
|
scope = anno.getanno(node, annos.NodeAnno.ARGS_AND_BODY_SCOPE)
|
|
closure_types = anno.getanno(node, anno.Static.CLOSURE_TYPES, {})
|
|
|
|
analyzer = Analyzer(subgraph, self.resolver, self.ctx.info.namespace, scope,
|
|
closure_types)
|
|
analyzer.visit_forward()
|
|
|
|
# Recursively process any remaining subfunctions.
|
|
node.body = self.visit_block(node.body)
|
|
|
|
return node
|
|
|
|
|
|
def resolve(node, source_info, graphs, resolver):
|
|
"""Performs type inference.
|
|
|
|
Args:
|
|
node: ast.AST
|
|
source_info: transformer.SourceInfo
|
|
graphs: Dict[ast.FunctionDef, cfg.Graph]
|
|
resolver: Resolver
|
|
|
|
Returns:
|
|
ast.AST
|
|
"""
|
|
visitor = FunctionVisitor(source_info, graphs, resolver)
|
|
node = visitor.visit(node)
|
|
return node
|