chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
License Check / license-check (push) Failing after 2s

This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
+416
View File
@@ -0,0 +1,416 @@
load("@rules_shell//shell:sh_test.bzl", "sh_test")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
py_library(
name = "templates",
srcs = ["templates.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":anno",
":ast_util",
":parser",
":qual_names",
"@pypi//gast",
],
)
py_library(
name = "transpiler",
srcs = ["transpiler.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":cache",
":inspect_utils",
":loader",
":naming",
":origin_info",
":parser",
":templates",
":transformer",
"//tensorflow/python/autograph/utils:ag_logging",
"@pypi//gast",
],
)
py_library(
name = "ast_util",
srcs = ["ast_util.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":anno",
":parser",
":qual_names",
"@pypi//gast",
],
)
py_library(
name = "loader",
srcs = ["loader.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":origin_info",
":parser",
],
)
py_library(
name = "gast_util",
srcs = ["gast_util.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = ["@pypi//gast"],
)
py_library(
name = "__init__",
srcs = ["__init__.py"],
strict_deps = True,
visibility = ["//visibility:public"],
)
py_library(
name = "parser",
srcs = ["parser.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":errors",
":inspect_utils",
"//tensorflow/python/util:tf_inspect",
"@pypi//gast",
],
)
py_library(
name = "naming",
srcs = ["naming.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [":qual_names"],
)
py_library(
name = "inspect_utils",
srcs = ["inspect_utils.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = ["//tensorflow/python/util:tf_inspect"],
)
py_library(
name = "origin_info",
srcs = ["origin_info.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":anno",
":ast_util",
":parser",
":pretty_printer",
"//tensorflow/python/util:tf_inspect",
"@pypi//gast",
],
)
py_library(
name = "anno",
srcs = ["anno.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = ["@pypi//gast"],
)
py_library(
name = "errors",
srcs = ["errors.py"],
strict_deps = True,
visibility = ["//visibility:public"],
)
py_library(
name = "transformer",
srcs = ["transformer.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":anno",
":parser",
":pretty_printer",
":templates",
"@pypi//gast",
],
)
py_library(
name = "qual_names",
srcs = ["qual_names.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":anno",
":parser",
"@pypi//gast",
],
)
py_library(
name = "cfg",
srcs = ["cfg.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":anno",
"@pypi//gast",
],
)
py_library(
name = "error_utils",
srcs = ["error_utils.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":origin_info",
"//tensorflow/python/util:traceback_utils",
],
)
py_library(
name = "cache",
srcs = ["cache.py"],
strict_deps = True,
visibility = ["//visibility:public"],
)
py_library(
name = "pretty_printer",
srcs = ["pretty_printer.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"@pypi//gast",
"@pypi//termcolor",
],
)
py_test(
name = "anno_test",
srcs = ["anno_test.py"],
strict_deps = True,
deps = [
":anno",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "ast_util_test",
srcs = ["ast_util_test.py"],
strict_deps = True,
deps = [
":anno",
":ast_util",
":loader",
":parser",
":pretty_printer",
":qual_names",
"@pypi//gast",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "cache_test",
srcs = ["cache_test.py"],
strict_deps = True,
deps = [
":cache",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "cfg_test",
srcs = ["cfg_test.py"],
strict_deps = True,
deps = [
":cfg",
":parser",
"@pypi//gast",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "loader_test",
srcs = ["loader_test.py"],
strict_deps = True,
deps = [
":ast_util",
":loader",
":parser",
":pretty_printer",
"@pypi//gast",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:tf_inspect",
],
)
py_test(
name = "error_utils_test",
srcs = ["error_utils_test.py"],
strict_deps = True,
deps = [
":error_utils",
":origin_info",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "inspect_utils_test",
srcs = ["inspect_utils_test.py"],
strict_deps = True,
deps = [
":inspect_utils",
#internal proto upb dep
"//tensorflow/python/autograph/pyct/testing:basic_definitions",
"//tensorflow/python/autograph/pyct/testing:decorators",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/lib:__init__",
"//tensorflow/python/platform:client_testlib",
],
)
sh_test(
name = "inspect_utils_test_par",
srcs = ["inspect_utils_test.sh"],
tags = ["no_oss"],
)
py_test(
name = "naming_test",
srcs = ["naming_test.py"],
strict_deps = True,
deps = [
":naming",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "origin_info_test",
srcs = ["origin_info_test.py"],
strict_deps = True,
deps = [
":anno",
":inspect_utils",
":origin_info",
":parser",
#internal proto upb dep
"//tensorflow/python/autograph/pyct/testing:basic_definitions",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:tf_inspect",
],
)
py_test(
name = "parser_test",
srcs = ["parser_test.py"],
strict_deps = True,
deps = [
":ast_util",
":errors",
":parser",
":pretty_printer",
"@pypi//gast",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "pretty_printer_test",
srcs = ["pretty_printer_test.py"],
strict_deps = True,
deps = [
":pretty_printer",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "qual_names_test",
srcs = ["qual_names_test.py"],
strict_deps = True,
deps = [
":anno",
":parser",
":qual_names",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "templates_test",
srcs = ["templates_test.py"],
strict_deps = True,
deps = [
":loader",
":parser",
":qual_names",
":templates",
"@absl_py//absl/testing:parameterized",
"@pypi//gast",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "transformer_test",
srcs = ["transformer_test.py"],
strict_deps = True,
deps = [
":anno",
":origin_info",
":parser",
":transformer",
"@pypi//gast",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "transpiler_test",
srcs = ["transpiler_test.py"],
strict_deps = True,
deps = [
":transformer",
":transpiler",
"@pypi//gast",
#internal proto upb dep
"//tensorflow/python/platform:client_testlib",
],
)
@@ -0,0 +1,16 @@
# Copyright 2017 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.
# ==============================================================================
"""Python source code transformation library."""
+174
View File
@@ -0,0 +1,174 @@
# Copyright 2017 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.
# ==============================================================================
"""AST node annotation support.
Adapted from Tangent.
"""
import enum
# pylint:disable=g-bad-import-order
import gast
# pylint:enable=g-bad-import-order
# TODO(mdan): Shorten the names.
# These names are heavily used, and anno.blaa
# TODO(mdan): Replace the attr-dict mechanism with a more typed solution.
class NoValue(enum.Enum):
"""Base class for different types of AST annotations."""
def of(self, node, default=None):
return getanno(node, self, default=default)
def add_to(self, node, value):
setanno(node, self, value)
def exists(self, node):
return hasanno(node, self)
def __repr__(self):
return str(self.name)
class Basic(NoValue):
"""Container for basic annotation keys.
The enum values are used strictly for documentation purposes.
"""
QN = 'Qualified name, as it appeared in the code. See qual_names.py.'
SKIP_PROCESSING = (
'This node should be preserved as is and not processed any further.')
INDENT_BLOCK_REMAINDER = (
'When a node is annotated with this, the remainder of the block should'
' be indented below it. The annotation contains a tuple'
' (new_body, name_map), where `new_body` is the new indented block and'
' `name_map` allows renaming symbols.')
ORIGIN = ('Information about the source code that converted code originated'
' from. See origin_information.py.')
DIRECTIVES = ('User directives associated with a statement or a variable.'
' Typically, they affect the immediately-enclosing statement.')
EXTRA_LOOP_TEST = (
'A special annotation containing additional test code to be executed in'
' for loops.')
class Static(NoValue):
"""Container for static analysis annotation keys.
The enum values are used strictly for documentation purposes.
"""
# Symbols
# These flags are boolean.
IS_PARAM = 'Symbol is a parameter to the function being analyzed.'
# Scopes
# Scopes are represented by objects of type activity.Scope.
SCOPE = 'The scope for the annotated node. See activity.py.'
# TODO(mdan): Drop these in favor of accessing the child's SCOPE.
ARGS_SCOPE = 'The scope for the argument list of a function call.'
COND_SCOPE = 'The scope for the test node of a conditional statement.'
BODY_SCOPE = (
'The scope for the main body of a statement (True branch for if '
'statements, main body for loops).')
ORELSE_SCOPE = (
'The scope for the orelse body of a statement (False branch for if '
'statements, orelse body for loops).')
# Static analysis annotations.
DEFINITIONS = (
'Reaching definition information. See reaching_definitions.py.')
ORIG_DEFINITIONS = (
'The value of DEFINITIONS that applied to the original code before any'
' conversion.')
DEFINED_FNS_IN = (
'Local function definitions that may exist when exiting the node. See'
' reaching_fndefs.py')
DEFINED_VARS_IN = (
'Symbols defined when entering the node. See reaching_definitions.py.')
LIVE_VARS_OUT = ('Symbols live when exiting the node. See liveness.py.')
LIVE_VARS_IN = ('Symbols live when entering the node. See liveness.py.')
TYPES = 'Static type information. See type_inference.py.'
CLOSURE_TYPES = 'Types of closure symbols at each detected call site.'
VALUE = 'Static value information. See type_inference.py.'
FAIL = object()
def keys(node, field_name='___pyct_anno'):
if not hasattr(node, field_name):
return frozenset()
return frozenset(getattr(node, field_name).keys())
def getanno(node, key, default=FAIL, field_name='___pyct_anno'):
if (default is FAIL or (hasattr(node, field_name) and
(key in getattr(node, field_name)))):
return getattr(node, field_name)[key]
return default
def hasanno(node, key, field_name='___pyct_anno'):
return hasattr(node, field_name) and key in getattr(node, field_name)
def setanno(node, key, value, field_name='___pyct_anno'):
annotations = getattr(node, field_name, {})
setattr(node, field_name, annotations)
annotations[key] = value
# So that the annotations survive gast_to_ast() and ast_to_gast()
if field_name not in node._fields:
node._fields += (field_name,)
def delanno(node, key, field_name='___pyct_anno'):
annotations = getattr(node, field_name)
del annotations[key]
if not annotations:
delattr(node, field_name)
node._fields = tuple(f for f in node._fields if f != field_name)
def copyanno(from_node, to_node, key, field_name='___pyct_anno'):
if hasanno(from_node, key, field_name=field_name):
setanno(
to_node,
key,
getanno(from_node, key, field_name=field_name),
field_name=field_name)
def dup(node, copy_map, field_name='___pyct_anno'):
"""Recursively copies annotations in an AST tree.
Args:
node: ast.AST
copy_map: Dict[Hashable, Hashable], maps a source anno key to a destination
key. All annotations with the source key will be copied to identical
annotations with the destination key.
field_name: str
"""
for n in gast.walk(node):
for k in copy_map:
if hasanno(n, k, field_name):
setanno(n, copy_map[k], getanno(n, k, field_name), field_name)
@@ -0,0 +1,80 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for anno module."""
import ast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.platform import test
# TODO(mdan): Consider strong types instead of primitives.
class AnnoTest(test.TestCase):
def test_basic(self):
node = ast.Name()
self.assertEqual(anno.keys(node), set())
self.assertFalse(anno.hasanno(node, 'foo'))
with self.assertRaises(AttributeError):
anno.getanno(node, 'foo')
anno.setanno(node, 'foo', 3)
self.assertEqual(anno.keys(node), {'foo'})
self.assertTrue(anno.hasanno(node, 'foo'))
self.assertEqual(anno.getanno(node, 'foo'), 3)
self.assertEqual(anno.getanno(node, 'bar', default=7), 7)
anno.delanno(node, 'foo')
self.assertEqual(anno.keys(node), set())
self.assertFalse(anno.hasanno(node, 'foo'))
with self.assertRaises(AttributeError):
anno.getanno(node, 'foo')
self.assertIsNone(anno.getanno(node, 'foo', default=None))
def test_copy(self):
node_1 = ast.Name()
anno.setanno(node_1, 'foo', 3)
node_2 = ast.Name()
anno.copyanno(node_1, node_2, 'foo')
anno.copyanno(node_1, node_2, 'bar')
self.assertTrue(anno.hasanno(node_2, 'foo'))
self.assertFalse(anno.hasanno(node_2, 'bar'))
def test_duplicate(self):
node = ast.If(
test=ast.Constant(value=1),
body=[ast.Expr(ast.Name('bar', ast.Load()))],
orelse=[])
anno.setanno(node, 'spam', 1)
anno.setanno(node, 'ham', 1)
anno.setanno(node.body[0], 'ham', 1)
anno.dup(node, {'spam': 'eggs'})
self.assertTrue(anno.hasanno(node, 'spam'))
self.assertTrue(anno.hasanno(node, 'ham'))
self.assertTrue(anno.hasanno(node, 'eggs'))
self.assertFalse(anno.hasanno(node.body[0], 'eggs'))
if __name__ == '__main__':
test.main()
@@ -0,0 +1,344 @@
# Copyright 2017 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.
# ==============================================================================
"""AST manipulation utilities."""
import ast
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names
class CleanCopier(object):
"""NodeTransformer-like visitor that copies an AST."""
def __init__(self, preserve_annos):
super(CleanCopier, self).__init__()
self.preserve_annos = preserve_annos
def copy(self, node):
"""Returns a deep copy of node (excluding some fields, see copy_clean)."""
if isinstance(node, list):
return [self.copy(n) for n in node]
elif isinstance(node, tuple):
return tuple(self.copy(n) for n in node)
elif not isinstance(node, (gast.AST, ast.AST)):
# Assuming everything that's not an AST, list or tuple is a value type
# and may simply be assigned.
return node
assert isinstance(node, (gast.AST, ast.AST))
new_fields = {}
for f in node._fields:
if not f.startswith('__') and hasattr(node, f):
new_fields[f] = self.copy(getattr(node, f))
new_node = type(node)(**new_fields)
if self.preserve_annos:
for k in self.preserve_annos:
anno.copyanno(node, new_node, k)
return new_node
def copy_clean(node, preserve_annos=None):
"""Creates a deep copy of an AST.
The copy will not include fields that are prefixed by '__', with the
exception of user-specified annotations.
Args:
node: ast.AST
preserve_annos: Optional[Set[Hashable]], annotation keys to include in the
copy
Returns:
ast.AST
"""
return CleanCopier(preserve_annos).copy(node)
class SymbolRenamer(gast.NodeTransformer):
"""Transformer that can rename symbols to a simple names."""
def __init__(self, name_map):
self.name_map = name_map
def _process_name_node(self, node):
qn = anno.getanno(node, anno.Basic.QN)
if qn in self.name_map:
new_node = gast.Name(
str(self.name_map[qn]),
ctx=node.ctx,
annotation=None,
type_comment=None)
# All annotations get carried over.
for k in anno.keys(node):
anno.copyanno(node, new_node, k)
return new_node
return self.generic_visit(node)
def _process_list_of_strings(self, names):
for i in range(len(names)):
qn = qual_names.QN(names[i])
if qn in self.name_map:
names[i] = str(self.name_map[qn])
return names
def visit_Nonlocal(self, node):
node.names = self._process_list_of_strings(node.names)
return node
def visit_Global(self, node):
node.names = self._process_list_of_strings(node.names)
return node
def visit_Name(self, node):
return self._process_name_node(node)
def visit_Attribute(self, node):
if anno.hasanno(node, anno.Basic.QN):
return self._process_name_node(node)
# Renaming attributes is not supported.
return self.generic_visit(node)
def visit_FunctionDef(self, node):
qn = qual_names.QN(node.name)
if qn in self.name_map:
node.name = str(self.name_map[qn])
return self.generic_visit(node)
def rename_symbols(node, name_map):
"""Renames symbols in an AST. Requires qual_names annotations."""
renamer = SymbolRenamer(name_map)
if isinstance(node, list):
return [renamer.visit(n) for n in node]
elif isinstance(node, tuple):
return tuple(renamer.visit(n) for n in node)
return renamer.visit(node)
def keywords_to_dict(keywords):
"""Converts a list of ast.keyword objects to a dict."""
keys = []
values = []
for kw in keywords:
keys.append(gast.Constant(kw.arg, kind=None))
values.append(kw.value)
return gast.Dict(keys=keys, values=values)
class PatternMatcher(gast.NodeVisitor):
"""Matches a node against a pattern represented by a node."""
def __init__(self, pattern):
self.pattern = pattern
self.pattern_stack = []
self.matches = True
def compare_and_visit(self, node, pattern):
self.pattern_stack.append(self.pattern)
self.pattern = pattern
self.generic_visit(node)
self.pattern = self.pattern_stack.pop()
def no_match(self):
self.matches = False
return False
def is_wildcard(self, p):
if isinstance(p, (list, tuple)) and len(p) == 1:
p, = p
if isinstance(p, gast.Name) and p.id == '_':
return True
if p == '_':
return True
return False
def generic_visit(self, node):
if not self.matches:
return
pattern = self.pattern
for f in node._fields:
if f.startswith('__'):
continue
if not hasattr(node, f):
if hasattr(pattern, f) and getattr(pattern, f):
return self.no_match()
else:
continue
if not hasattr(pattern, f):
return self.no_match()
v = getattr(node, f)
p = getattr(pattern, f)
if self.is_wildcard(p):
continue
if isinstance(v, (list, tuple)):
if not isinstance(p, (list, tuple)) or len(v) != len(p):
return self.no_match()
for v_item, p_item in zip(v, p):
self.compare_and_visit(v_item, p_item)
elif isinstance(v, (gast.AST, ast.AST)):
if not isinstance(v, type(p)) and not isinstance(p, type(v)):
return self.no_match()
self.compare_and_visit(v, p)
else:
# Assume everything else is a value type.
if v != p:
return self.no_match()
def matches(node, pattern):
"""Basic pattern matcher for AST.
The pattern may contain wildcards represented by the symbol '_'. A node
matches a pattern if for every node in the tree, either there is a node of
the same type in pattern, or a Name node with id='_'.
Args:
node: ast.AST
pattern: ast.AST
Returns:
bool
"""
if isinstance(pattern, str):
pattern = parser.parse_str(pattern)
matcher = PatternMatcher(pattern)
matcher.visit(node)
return matcher.matches
# TODO(mdan): Once we have error tracing, we may be able to just go to SSA.
def apply_to_single_assignments(targets, values, apply_fn):
"""Applies a function to each individual assignment.
This function can process a possibly-unpacked (e.g. a, b = c, d) assignment.
It tries to break down the unpacking if possible. In effect, it has the same
effect as passing the assigned values in SSA form to apply_fn.
Examples:
The following will result in apply_fn(a, c), apply_fn(b, d):
a, b = c, d
The following will result in apply_fn(a, c[0]), apply_fn(b, c[1]):
a, b = c
The following will result in apply_fn(a, (b, c)):
a = b, c
It uses the visitor pattern to allow subclasses to process single
assignments individually.
Args:
targets: Union[List[ast.AST, ...], Tuple[ast.AST, ...], ast.AST, should be
used with the targets field of an ast.Assign node
values: ast.AST
apply_fn: Callable[[ast.AST, ast.AST], None], called with the
respective nodes of each single assignment
"""
if not isinstance(targets, (list, tuple)):
targets = (targets,)
for target in targets:
if isinstance(target, (gast.Tuple, gast.List)):
for i in range(len(target.elts)):
target_el = target.elts[i]
if isinstance(values, (gast.Tuple, gast.List)):
value_el = values.elts[i]
else:
idx = parser.parse_expression(str(i))
value_el = gast.Subscript(values, idx, ctx=gast.Load())
apply_to_single_assignments(target_el, value_el, apply_fn)
else:
apply_fn(target, values)
def parallel_walk(node, other):
"""Walks two ASTs in parallel.
The two trees must have identical structure.
Args:
node: Union[ast.AST, Iterable[ast.AST]]
other: Union[ast.AST, Iterable[ast.AST]]
Yields:
Tuple[ast.AST, ast.AST]
Raises:
ValueError: if the two trees don't have identical structure.
"""
if isinstance(node, (list, tuple)):
node_stack = list(node)
else:
node_stack = [node]
if isinstance(other, (list, tuple)):
other_stack = list(other)
else:
other_stack = [other]
while node_stack and other_stack:
assert len(node_stack) == len(other_stack)
n = node_stack.pop()
o = other_stack.pop()
if ((not isinstance(n, (ast.AST, gast.AST, str)) and n is not None) or
(not isinstance(o, (ast.AST, gast.AST, str)) and n is not None) or
n.__class__.__name__ != o.__class__.__name__):
raise ValueError('inconsistent nodes: {} ({}) and {} ({})'.format(
n, n.__class__.__name__, o, o.__class__.__name__))
yield n, o
if isinstance(n, str):
assert isinstance(o, str), 'The check above should have ensured this'
continue
if n is None:
assert o is None, 'The check above should have ensured this'
continue
for f in n._fields:
n_child = getattr(n, f, None)
o_child = getattr(o, f, None)
if f.startswith('__') or n_child is None or o_child is None:
continue
if isinstance(n_child, (list, tuple)):
if (not isinstance(o_child, (list, tuple)) or
len(n_child) != len(o_child)):
raise ValueError(
'inconsistent values for field {}: {} and {}'.format(
f, n_child, o_child))
node_stack.extend(n_child)
other_stack.extend(o_child)
elif isinstance(n_child, (gast.AST, ast.AST)):
node_stack.append(n_child)
other_stack.append(o_child)
elif n_child != o_child:
raise ValueError(
'inconsistent values for field {}: {} and {}'.format(
f, n_child, o_child))
@@ -0,0 +1,246 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for ast_util module."""
import ast
import collections
import textwrap
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import ast_util
from tensorflow.python.autograph.pyct import loader
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import pretty_printer
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.platform import test
class AstUtilTest(test.TestCase):
def assertAstMatches(self, actual_node, expected_node_src):
expected_node = gast.parse('({})'.format(expected_node_src)).body[0]
msg = 'AST did not match expected:\n{}\nActual:\n{}'.format(
pretty_printer.fmt(expected_node),
pretty_printer.fmt(actual_node))
self.assertTrue(ast_util.matches(actual_node, expected_node), msg)
def setUp(self):
super(AstUtilTest, self).setUp()
self._invocation_counts = collections.defaultdict(lambda: 0)
def test_rename_symbols_basic(self):
node = parser.parse('a + b')
node = qual_names.resolve(node)
node = ast_util.rename_symbols(
node, {qual_names.QN('a'): qual_names.QN('renamed_a')})
source = parser.unparse(node, include_encoding_marker=False)
expected_node_src = 'renamed_a + b'
self.assertIsInstance(node.value.left.id, str)
self.assertAstMatches(node, source)
self.assertAstMatches(node, expected_node_src)
def test_rename_symbols_attributes(self):
node = parser.parse('b.c = b.c.d')
node = qual_names.resolve(node)
node = ast_util.rename_symbols(
node, {qual_names.from_str('b.c'): qual_names.QN('renamed_b_c')})
source = parser.unparse(node, include_encoding_marker=False)
self.assertEqual(source.strip(), 'renamed_b_c = renamed_b_c.d')
def test_rename_symbols_nonlocal(self):
node = parser.parse('nonlocal a, b, c')
node = qual_names.resolve(node)
node = ast_util.rename_symbols(
node, {qual_names.from_str('b'): qual_names.QN('renamed_b')})
source = parser.unparse(node, include_encoding_marker=False)
self.assertEqual(source.strip(), 'nonlocal a, renamed_b, c')
def test_rename_symbols_global(self):
node = parser.parse('global a, b, c')
node = qual_names.resolve(node)
node = ast_util.rename_symbols(
node, {qual_names.from_str('b'): qual_names.QN('renamed_b')})
source = parser.unparse(node, include_encoding_marker=False)
self.assertEqual(source.strip(), 'global a, renamed_b, c')
def test_rename_symbols_annotations(self):
node = parser.parse('a[i]')
node = qual_names.resolve(node)
anno.setanno(node, 'foo', 'bar')
orig_anno = anno.getanno(node, 'foo')
node = ast_util.rename_symbols(node,
{qual_names.QN('a'): qual_names.QN('b')})
self.assertIs(anno.getanno(node, 'foo'), orig_anno)
def test_rename_symbols_function(self):
node = parser.parse('def f():\n pass')
node = ast_util.rename_symbols(node,
{qual_names.QN('f'): qual_names.QN('f1')})
source = parser.unparse(node, include_encoding_marker=False)
self.assertEqual(source.strip(), 'def f1():\n pass')
def test_copy_clean(self):
node = parser.parse(
textwrap.dedent("""
def f(a):
return a + 1
"""))
setattr(node, '__foo', 'bar')
new_node = ast_util.copy_clean(node)
self.assertIsNot(new_node, node)
self.assertFalse(hasattr(new_node, '__foo'))
def test_copy_clean_preserves_annotations(self):
node = parser.parse(
textwrap.dedent("""
def f(a):
return a + 1
"""))
anno.setanno(node, 'foo', 'bar')
anno.setanno(node, 'baz', 1)
new_node = ast_util.copy_clean(node, preserve_annos={'foo'})
self.assertEqual(anno.getanno(new_node, 'foo'), 'bar')
self.assertFalse(anno.hasanno(new_node, 'baz'))
def test_keywords_to_dict(self):
keywords = parser.parse_expression('f(a=b, c=1, d=\'e\')').keywords
d = ast_util.keywords_to_dict(keywords)
# Make sure we generate a usable dict node by attaching it to a variable and
# compiling everything.
node = parser.parse('def f(b): pass')
node.body.append(ast.Return(d))
result, _, _ = loader.load_ast(node)
self.assertDictEqual(result.f(3), {'a': 3, 'c': 1, 'd': 'e'})
def assertMatch(self, target_str, pattern_str):
node = parser.parse_expression(target_str)
pattern = parser.parse_expression(pattern_str)
self.assertTrue(ast_util.matches(node, pattern))
def assertNoMatch(self, target_str, pattern_str):
node = parser.parse_expression(target_str)
pattern = parser.parse_expression(pattern_str)
self.assertFalse(ast_util.matches(node, pattern))
def test_matches_symbols(self):
self.assertMatch('foo', '_')
self.assertNoMatch('foo()', '_')
self.assertMatch('foo + bar', 'foo + _')
self.assertNoMatch('bar + bar', 'foo + _')
self.assertNoMatch('foo - bar', 'foo + _')
def test_matches_function_args(self):
self.assertMatch('super(Foo, self).__init__(arg1, arg2)',
'super(_).__init__(_)')
self.assertMatch('super().__init__()', 'super(_).__init__(_)')
self.assertNoMatch('super(Foo, self).bar(arg1, arg2)',
'super(_).__init__(_)')
self.assertMatch('super(Foo, self).__init__()', 'super(Foo, _).__init__(_)')
self.assertNoMatch('super(Foo, self).__init__()',
'super(Bar, _).__init__(_)')
def _mock_apply_fn(self, target, source):
target = parser.unparse(target, include_encoding_marker=False)
source = parser.unparse(source, include_encoding_marker=False)
self._invocation_counts[(target.strip(), source.strip())] += 1
def test_apply_to_single_assignments_dynamic_unpack(self):
node = parser.parse('a, b, c = d')
ast_util.apply_to_single_assignments(node.targets, node.value,
self._mock_apply_fn)
self.assertDictEqual(self._invocation_counts, {
('a', 'd[0]'): 1,
('b', 'd[1]'): 1,
('c', 'd[2]'): 1,
})
def test_apply_to_single_assignments_static_unpack(self):
node = parser.parse('a, b, c = d, e, f')
ast_util.apply_to_single_assignments(node.targets, node.value,
self._mock_apply_fn)
self.assertDictEqual(self._invocation_counts, {
('a', 'd'): 1,
('b', 'e'): 1,
('c', 'f'): 1,
})
def test_parallel_walk(self):
src = """
def f(a):
return a + 1
"""
node = parser.parse(textwrap.dedent(src))
for child_a, child_b in ast_util.parallel_walk(node, node):
self.assertEqual(child_a, child_b)
def test_parallel_walk_string_leaves(self):
src = """
def f(a):
global g
"""
node = parser.parse(textwrap.dedent(src))
for child_a, child_b in ast_util.parallel_walk(node, node):
self.assertEqual(child_a, child_b)
def test_parallel_walk_inconsistent_trees(self):
node_1 = parser.parse(
textwrap.dedent("""
def f(a):
return a + 1
"""))
node_2 = parser.parse(
textwrap.dedent("""
def f(a):
return a + (a * 2)
"""))
node_3 = parser.parse(
textwrap.dedent("""
def f(a):
return a + 2
"""))
with self.assertRaises(ValueError):
for _ in ast_util.parallel_walk(node_1, node_2):
pass
# There is not particular reason to reject trees that differ only in the
# value of a constant.
# TODO(mdan): This should probably be allowed.
with self.assertRaises(ValueError):
for _ in ast_util.parallel_walk(node_1, node_3):
pass
def assertLambdaNodes(self, matching_nodes, expected_bodies):
self.assertEqual(len(matching_nodes), len(expected_bodies))
for node in matching_nodes:
self.assertIsInstance(node, gast.Lambda)
self.assertIn(
parser.unparse(node.body, include_encoding_marker=False).strip(),
expected_bodies)
if __name__ == '__main__':
test.main()
+93
View File
@@ -0,0 +1,93 @@
# Copyright 2016 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.
# ==============================================================================
"""Caching utilities."""
import inspect
import weakref
# TODO(mdan): Add a garbage collection hook for cleaning up modules.
class _TransformedFnCache(object):
"""Generic hierarchical cache for transformed functions.
The keys are soft references (i.e. they are discarded when the key is
destroyed) created from the source function by `_get_key`. The subkeys are
strong references and can be any value. Typically they identify different
kinds of transformation.
"""
__slots__ = ('_cache',)
def __init__(self):
self._cache = weakref.WeakKeyDictionary()
def _get_key(self, entity):
raise NotImplementedError('subclasses must override')
def has(self, entity, subkey):
key = self._get_key(entity)
parent = self._cache.get(key, None)
if parent is None:
return False
return subkey in parent
def __getitem__(self, entity):
key = self._get_key(entity)
parent = self._cache.get(key, None)
if parent is None:
# The bucket is initialized to support this usage:
# cache[key][subkey] = value
self._cache[key] = parent = {}
return parent
def __len__(self):
return len(self._cache)
class CodeObjectCache(_TransformedFnCache):
"""A function cache based on code objects.
Code objects are good proxies for the source code of a function.
This cache efficiently handles functions that share code objects, such as
functions defined in a loop, bound methods, etc.
The cache falls back to the function object, if it doesn't have a code object.
"""
def _get_key(self, entity):
if hasattr(entity, '__code__'):
return entity.__code__
else:
return entity
class UnboundInstanceCache(_TransformedFnCache):
"""A function cache based on unbound function objects.
Using the function for the cache key allows efficient handling of object
methods.
Unlike the _CodeObjectCache, this discriminates between different functions
even if they have the same code. This is needed for decorators that may
masquerade as another function.
"""
def _get_key(self, entity):
if inspect.ismethod(entity):
return entity.__func__
return entity
@@ -0,0 +1,75 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for cache module."""
from tensorflow.python.autograph.pyct import cache
from tensorflow.python.platform import test
class CacheTest(test.TestCase):
def test_code_object_cache(self):
def factory(x):
def test_fn():
return x + 1
return test_fn
c = cache.CodeObjectCache()
f1 = factory(1)
dummy = object()
c[f1][1] = dummy
self.assertTrue(c.has(f1, 1))
self.assertFalse(c.has(f1, 2))
self.assertIs(c[f1][1], dummy)
self.assertEqual(len(c), 1)
f2 = factory(2)
self.assertTrue(c.has(f2, 1))
self.assertIs(c[f2][1], dummy)
self.assertEqual(len(c), 1)
def test_unbound_instance_cache(self):
class TestClass(object):
def method(self):
pass
c = cache.UnboundInstanceCache()
o1 = TestClass()
dummy = object()
c[o1.method][1] = dummy
self.assertTrue(c.has(o1.method, 1))
self.assertFalse(c.has(o1.method, 2))
self.assertIs(c[o1.method][1], dummy)
self.assertEqual(len(c), 1)
o2 = TestClass()
self.assertTrue(c.has(o2.method, 1))
self.assertIs(c[o2.method][1], dummy)
self.assertEqual(len(c), 1)
if __name__ == '__main__':
test.main()
+975
View File
@@ -0,0 +1,975 @@
# Copyright 2017 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.
# ==============================================================================
"""Control flow graph (CFG) structure for Python AST representation.
The CFG is a digraph with edges representing valid control flow. Each
node is associated with exactly one AST node, but not all AST nodes may have
a corresponding CFG counterpart.
Once built, the CFG itself is immutable, but the values it holds need not be;
they are usually annotated with information extracted by walking the graph.
Tip: Use `Graph.as_dot` to visualize the CFG using any DOT viewer.
Note: the CFG tries to include all code paths that MAY be taken, with a single
notable exception:
* function calls do not generate edges corresponding to exceptions they may
raise (i.e. a function call in the middle of a block does not return or jump
to any except or finally block)
TODO(mdan): Consider adding the edges above. They'd only add ~O(n) edges.
TODO(mdan): Alternatively, consider adding an edge from try to all its excepts.
"""
# TODO(mdan): The notion of 'statements' below is inaccurate.
# They should rather be called 'block statements', because they include
# statements that may have a body, e.g. if and while.
import ast
import collections
import enum
import weakref
import gast
from tensorflow.python.autograph.pyct import anno
class Node(object):
"""A node in the CFG.
Although new instances of this class are mutable, the objects that a user
finds in the CFG are typically not.
The nodes represent edges in the CFG graph, and maintain pointers to allow
efficient walking in both forward and reverse order. The following property
holds for all nodes: "child in node.next" iff "node in child.prev".
Attributes:
next: FrozenSet[Node, ...], the nodes that follow this node, in control flow
order
prev: FrozenSet[Node, ...], the nodes that precede this node, in reverse
control flow order
ast_node: ast.AST, the AST node corresponding to this CFG node
"""
def __init__(self, next_, prev, ast_node):
self.next = next_
self.prev = prev
self.ast_node = ast_node
def freeze(self):
self.next = frozenset(self.next)
# Assumption: All CFG nodes have identical life spans, because the graph
# owns them. Nodes should never be used outside the context of an existing
# graph.
self.prev = weakref.WeakSet(self.prev)
def __repr__(self):
if isinstance(self.ast_node, gast.FunctionDef):
return 'def %s' % self.ast_node.name
elif isinstance(self.ast_node, gast.ClassDef):
return 'class %s' % self.ast_node.name
elif isinstance(self.ast_node, gast.withitem):
return ast.unparse(self.ast_node.context_expr).strip()
return ast.unparse(self.ast_node).strip()
class Graph(
collections.namedtuple(
'Graph',
['entry', 'exit', 'error', 'index', 'stmt_prev', 'stmt_next'])):
"""A Control Flow Graph.
The CFG maintains an index to allow looking up a CFG node by the AST node to
which it is associated. The index can also be enumerated in top-down, depth
first order.
Walking the graph in forward or reverse order is supported by double
parent-child links.
Note: the error nodes are not wired to their corresponding finally guards,
because these are shared, and wiring them would create a reverse path from
normal control flow into the error nodes, which we want to avoid.
The graph also maintains edges corresponding to higher level statements
like for-else loops. A node is considered successor of a statement if there
is an edge from a node that is lexically a child of that statement to a node
that is not. Statement predecessors are analogously defined.
Attributes:
entry: Node, the entry node
exit: FrozenSet[Node, ...], the exit nodes
error: FrozenSet[Node, ...], nodes that exit due to an explicitly raised
error (errors propagated from function calls are not accounted)
index: Dict[ast.Node, Node], mapping AST nodes to the respective CFG node
stmt_prev: Dict[ast.Node, FrozenSet[Node, ...]], mapping statement AST nodes
to their predecessor CFG nodes
stmt_next: Dict[ast.Node, FrozenSet[Node, ...]], mapping statement AST nodes
to their successor CFG nodes
"""
def __repr__(self):
return self.as_dot()
def as_dot(self):
"""Print CFG in DOT format."""
result = 'digraph CFG {\n'
for node in self.index.values():
result += ' %s [label="%s"];\n' % (id(node), node)
for node in self.index.values():
for next_ in node.next:
result += ' %s -> %s;\n' % (id(node), id(next_))
result += '}'
return result
class _WalkMode(enum.Enum):
FORWARD = 1
REVERSE = 2
# TODO(mdan): Rename to DataFlowAnalyzer.
# TODO(mdan): Consider specializations that use gen/kill/transfer abstractions.
class GraphVisitor(object):
"""Base class for a CFG visitors.
This implementation is not thread safe.
The visitor has some facilities to simplify dataflow analyses. In particular,
it allows revisiting the nodes at the decision of the subclass. This can be
used to visit the graph until the state reaches a fixed point.
For more details on dataflow analysis, see
https://www.seas.harvard.edu/courses/cs252/2011sp/slides/Lec02-Dataflow.pdf
Note: the literature generally suggests visiting successor nodes only when the
state of the current node changed, regardless of whether that successor has
ever been visited. This implementation visits every successor at least once.
Attributes:
graph: Graph
in_: Dict[Node, Any], stores node-keyed state during a visit
out: Dict[Node, Any], stores node-keyed state during a visit
"""
def __init__(self, graph):
self.graph = graph
self.reset()
def init_state(self, node):
"""State initialization function.
Optional to overload.
An in/out state slot will be created for each node in the graph. Subclasses
must overload this to control what that is initialized to.
Args:
node: Node
"""
raise NotImplementedError('Subclasses must implement this.')
# TODO(mdan): Rename to flow?
def visit_node(self, node):
"""Visitor function.
Args:
node: Node
Returns:
bool, whether the node should be revisited; subclasses can visit every
reachable node exactly once by always returning False
"""
raise NotImplementedError('Subclasses must implement this.')
def reset(self):
self.in_ = {
node: self.init_state(node) for node in self.graph.index.values()
}
self.out = {
node: self.init_state(node) for node in self.graph.index.values()
}
def can_ignore(self, node):
"""Returns True if the node can safely be assumed not to touch variables."""
ast_node = node.ast_node
if anno.hasanno(ast_node, anno.Basic.SKIP_PROCESSING):
return True
return isinstance(ast_node,
(gast.Break, gast.Continue, gast.Raise, gast.Pass))
def _visit_internal(self, mode):
"""Visits the CFG, breadth-first."""
assert mode in (_WalkMode.FORWARD, _WalkMode.REVERSE)
if mode == _WalkMode.FORWARD:
open_ = [self.graph.entry]
elif mode == _WalkMode.REVERSE:
open_ = list(self.graph.exit)
closed = set()
while open_:
node = open_.pop(0)
closed.add(node)
should_revisit = self.visit_node(node)
if mode == _WalkMode.FORWARD:
children = node.next
elif mode == _WalkMode.REVERSE:
children = node.prev
for next_ in children:
if should_revisit or next_ not in closed:
open_.append(next_)
def visit_forward(self):
self._visit_internal(_WalkMode.FORWARD)
def visit_reverse(self):
self._visit_internal(_WalkMode.REVERSE)
class GraphBuilder(object):
"""Builder that constructs a CFG from a given AST.
This GraphBuilder facilitates constructing the DAG that forms the CFG when
nodes
are supplied in lexical order (i.e., top-down, depth first). Under these
conditions, it supports building patterns found in typical structured
programs.
This builder ignores the flow generated by exceptions, which are assumed to
always be catastrophic and present purely for diagnostic purposes (e.g. to
print debug information). Statements like raise and try/catch sections are
allowed and will generate control flow edges, but ordinary statements are
assumed not to raise exceptions.
Finally sections are also correctly interleaved between break/continue/return
nodes and their subsequent statements.
Important concepts:
* nodes - nodes refer to CFG nodes; AST nodes are qualified explicitly
* leaf set - since the graph is constructed gradually, a leaf set maintains
the CFG nodes that will precede the node that the builder expects to
receive next; when an ordinary node is added, it is connected to the
existing leaves and it in turn becomes the new leaf
* jump nodes - nodes that should generate edges other than what
ordinary nodes would; these correspond to break, continue and return
statements
* sections - logical delimiters for subgraphs that require special
edges; there are various types of nodes, each admitting various
types of jump nodes; sections are identified by their corresponding AST
node
"""
# TODO(mdan): Perhaps detail this in a markdown doc.
# TODO(mdan): Add exception support.
def __init__(self, parent_ast_node):
self.reset()
self.parent = parent_ast_node
def reset(self):
"""Resets the state of this factory."""
self.head = None
self.errors = set()
self.node_index = {}
# TODO(mdan): Too many primitives. Use classes.
self.leaves = set()
# Note: This mechanism requires that nodes are added in lexical order (top
# to bottom, depth first).
self.active_stmts = set()
self.owners = {} # type: Set[any]
self.forward_edges = set() # type: Tuple[Node, Node] # (from, to)
self.finally_sections = {}
# Dict values represent (entry, exits)
self.finally_section_subgraphs = {
} # type: Dict[ast.AST, Tuple[Node, Set[Node]]]
# Whether the guard section can be reached from the statement that precedes
# it.
self.finally_section_has_direct_flow = {}
# Finally sections that await their first node.
self.pending_finally_sections = set()
# Exit jumps keyed by the section they affect.
self.exits = {}
# The entry of loop sections, keyed by the section.
self.section_entry = {}
# Continue jumps keyed by the section they affect.
self.continues = {}
# Raise jumps keyed by the except section guarding them.
self.raises = {}
# The entry of conditional sections, keyed by the section.
self.cond_entry = {}
# Lists of leaf nodes corresponding to each branch in the section.
self.cond_leaves = {}
def _connect_nodes(self, first, second):
"""Connects nodes to signify that control flows from first to second.
Args:
first: Union[Set[Node, ...], Node]
second: Node
"""
if isinstance(first, Node):
first.next.add(second)
second.prev.add(first)
self.forward_edges.add((first, second))
else:
for node in first:
self._connect_nodes(node, second)
def _add_new_node(self, ast_node):
"""Grows the graph by adding a CFG node following the current leaves."""
if ast_node in self.node_index:
raise ValueError('%s added twice' % ast_node)
# Assumption: All CFG nodes have identical life spans, because the graph
# owns them. Nodes should never be used outside the context of an existing
# graph.
node = Node(next_=set(), prev=weakref.WeakSet(), ast_node=ast_node)
self.node_index[ast_node] = node
self.owners[node] = frozenset(self.active_stmts)
if self.head is None:
self.head = node
for leaf in self.leaves:
self._connect_nodes(leaf, node)
# If any finally section awaits its first node, populate it.
for section_id in self.pending_finally_sections:
self.finally_section_subgraphs[section_id][0] = node
self.pending_finally_sections = set()
return node
def begin_statement(self, stmt):
"""Marks the beginning of a statement.
Args:
stmt: Hashable, a key by which the statement can be identified in the
CFG's stmt_prev and stmt_next attributes
"""
self.active_stmts.add(stmt)
def end_statement(self, stmt):
"""Marks the end of a statement.
Args:
stmt: Hashable, a key by which the statement can be identified in the
CFG's stmt_prev and stmt_next attributes; must match a key previously
passed to begin_statement.
"""
self.active_stmts.remove(stmt)
def add_ordinary_node(self, ast_node):
"""Grows the graph by adding an ordinary CFG node.
Ordinary nodes are followed by the next node, in lexical order, that is,
they become the new leaf set.
Args:
ast_node: ast.AST
Returns:
Node
"""
node = self._add_new_node(ast_node)
self.leaves = set((node,))
return node
def _add_jump_node(self, ast_node, guards):
"""Grows the graph by adding a jump node.
Jump nodes are added to the current leaf set, and the leaf set becomes
empty. If the jump node is the last in a cond section, then it may be added
back to the leaf set by a separate mechanism.
Args:
ast_node: ast.AST
guards: Tuple[ast.AST, ...], the finally sections active for this node
Returns:
Node
"""
node = self._add_new_node(ast_node)
self.leaves = set()
# The guards themselves may not yet be complete, and will be wired later.
self.finally_sections[node] = guards
return node
def _connect_jump_to_finally_sections(self, node):
"""Connects a jump node to the finally sections protecting it."""
cursor = set((node,))
if node not in self.finally_sections:
return cursor
for guard_section_id in self.finally_sections[node]:
guard_begin, guard_ends = self.finally_section_subgraphs[guard_section_id]
self._connect_nodes(cursor, guard_begin)
cursor = guard_ends
del self.finally_sections[node]
# TODO(mdan): Should garbage-collect finally_section_subgraphs.
return cursor
def add_exit_node(self, ast_node, section_id, guards):
"""Grows the graph by adding an exit node.
This node becomes an exit for the current section.
Args:
ast_node: ast.AST
section_id: Hashable, the node for which ast_node should be considered to
be an exit node
guards: Tuple[ast.AST, ...], the finally sections that guard ast_node
Returns:
Node
"""
node = self._add_jump_node(ast_node, guards)
self.exits[section_id].add(node)
return node
def add_continue_node(self, ast_node, section_id, guards):
"""Grows the graph by adding a reentry node.
This node causes control flow to go back to the loop section's entry.
Args:
ast_node: ast.AST
section_id: Hashable, the node for which ast_node should be considered to
be an exit node
guards: Tuple[ast.AST, ...], the finally sections that guard ast_node
"""
node = self._add_jump_node(ast_node, guards)
self.continues[section_id].add(node)
def connect_raise_node(self, node, except_guards):
"""Adds extra connection between a raise node and containing except guards.
The node is a graph node, not an ast node.
Args:
node: Node
except_guards: Tuple[ast.AST, ...], the except sections that guard node
"""
for guard in except_guards:
if guard in self.raises:
self.raises[guard].append(node)
else:
self.raises[guard] = [node]
def enter_section(self, section_id):
"""Enters a regular section.
Regular sections admit exit jumps, which end the section.
Args:
section_id: Hashable, the same node that will be used in calls to the
ast_node arg passed to add_exit_node
"""
assert section_id not in self.exits
self.exits[section_id] = set()
def exit_section(self, section_id):
"""Exits a regular section."""
# Exits are jump nodes, which may be protected.
for exit_ in self.exits[section_id]:
self.leaves |= self._connect_jump_to_finally_sections(exit_)
del self.exits[section_id]
def enter_loop_section(self, section_id, entry_node):
"""Enters a loop section.
Loop sections define an entry node. The end of the section always flows back
to the entry node. These admit continue jump nodes which also flow to the
entry node.
Args:
section_id: Hashable, the same node that will be used in calls to the
ast_node arg passed to add_continue_node
entry_node: ast.AST, the entry node into the loop (e.g. the test node for
while loops)
"""
assert section_id not in self.section_entry
assert section_id not in self.continues
self.continues[section_id] = set()
node = self.add_ordinary_node(entry_node)
self.section_entry[section_id] = node
def exit_loop_section(self, section_id):
"""Exits a loop section."""
self._connect_nodes(self.leaves, self.section_entry[section_id])
# continues are jump nodes, which may be protected.
for reentry in self.continues[section_id]:
guard_ends = self._connect_jump_to_finally_sections(reentry)
self._connect_nodes(guard_ends, self.section_entry[section_id])
# Loop nodes always loop back.
self.leaves = set((self.section_entry[section_id],))
del self.continues[section_id]
del self.section_entry[section_id]
def enter_cond_section(self, section_id):
"""Enters a conditional section.
Conditional sections define an entry node, and one or more branches.
Args:
section_id: Hashable, the same node that will be used in calls to the
section_id arg passed to new_cond_branch
"""
assert section_id not in self.cond_entry
assert section_id not in self.cond_leaves
self.cond_leaves[section_id] = []
def new_cond_branch(self, section_id):
"""Begins a new branch in a cond section."""
assert section_id in self.cond_leaves
if section_id in self.cond_entry:
# Subsequent splits move back to the split point, and memorize the
# current leaves.
self.cond_leaves[section_id].append(self.leaves)
self.leaves = self.cond_entry[section_id]
else:
# If this is the first time we split a section, just remember the split
# point.
self.cond_entry[section_id] = self.leaves
def exit_cond_section(self, section_id):
"""Exits a conditional section."""
for split in self.cond_leaves[section_id]:
self.leaves |= split
del self.cond_entry[section_id]
del self.cond_leaves[section_id]
def enter_except_section(self, section_id):
"""Enters an except section."""
if section_id in self.raises:
self.leaves.update(self.raises[section_id])
def enter_finally_section(self, section_id):
"""Enters a finally section."""
# TODO(mdan): This, not the caller, should track the active sections.
self.finally_section_subgraphs[section_id] = [None, None]
if self.leaves:
self.finally_section_has_direct_flow[section_id] = True
else:
self.finally_section_has_direct_flow[section_id] = False
self.pending_finally_sections.add(section_id)
def exit_finally_section(self, section_id):
"""Exits a finally section."""
assert section_id not in self.pending_finally_sections, 'Empty finally?'
self.finally_section_subgraphs[section_id][1] = self.leaves
# If the guard can only be reached by a jump, then it will not flow
# into the statement that follows it.
if not self.finally_section_has_direct_flow[section_id]:
self.leaves = set()
del self.finally_section_has_direct_flow[section_id]
def build(self):
"""Returns the CFG accumulated so far and resets the builder.
Returns:
Graph
"""
# Freeze the nodes.
for node in self.node_index.values():
node.freeze()
# Build the statement edges.
stmt_next = {}
stmt_prev = {}
for node in self.node_index.values():
for stmt in self.owners[node]:
if stmt not in stmt_prev:
stmt_prev[stmt] = set()
if stmt not in stmt_next:
stmt_next[stmt] = set()
for first, second in self.forward_edges:
stmts_exited = self.owners[first] - self.owners[second]
for stmt in stmts_exited:
stmt_next[stmt].add(second)
stmts_entered = self.owners[second] - self.owners[first]
for stmt in stmts_entered:
stmt_prev[stmt].add(first)
for stmt in stmt_next:
stmt_next[stmt] = frozenset(stmt_next[stmt])
for stmt in stmt_prev:
stmt_prev[stmt] = frozenset(stmt_prev[stmt])
# Construct the final graph object.
result = Graph(
entry=self.head,
exit=self.leaves,
error=self.errors,
index=self.node_index,
stmt_prev=stmt_prev,
stmt_next=stmt_next)
# Reset the state.
self.reset()
return result
class AstToCfg(gast.NodeVisitor):
"""Converts an AST to CFGs.
A separate CFG will be constructed for each function.
"""
def __init__(self):
super(AstToCfg, self).__init__()
self.builder_stack = []
self.builder = None
self.cfgs = {}
self.lexical_scopes = []
def _enter_lexical_scope(self, node):
self.lexical_scopes.append(node)
def _exit_lexical_scope(self, node):
leaving_node = self.lexical_scopes.pop()
assert node == leaving_node
def _get_enclosing_finally_scopes(self, stop_at):
included = []
for node in reversed(self.lexical_scopes):
if isinstance(node, gast.Try) and node.finalbody:
included.append(node)
if isinstance(node, stop_at):
return node, included
return None, included
def _get_enclosing_except_scopes(self, stop_at):
included = []
for node in reversed(self.lexical_scopes):
if isinstance(node, gast.Try) and node.handlers:
included.extend(node.handlers)
if isinstance(node, stop_at):
break
return included
def _process_basic_statement(self, node):
self.generic_visit(node)
self.builder.add_ordinary_node(node)
def _process_exit_statement(self,
node,
exits_nodes_of_type,
may_exit_via_except=False):
self.generic_visit(node)
# Note: this is safe because we process functions separately.
try_node, guards = self._get_enclosing_finally_scopes(exits_nodes_of_type)
assert try_node is not None, '{} that is not enclosed by any of {}'.format(
node, exits_nodes_of_type)
node = self.builder.add_exit_node(node, try_node, guards)
if may_exit_via_except:
except_guards = self._get_enclosing_except_scopes(exits_nodes_of_type)
self.builder.connect_raise_node(node, except_guards)
def _process_continue_statement(self, node, *loops_to_nodes_of_type):
# Note: this is safe because we process functions separately.
try_node, guards = self._get_enclosing_finally_scopes(
tuple(loops_to_nodes_of_type))
if try_node is None:
raise ValueError('%s that is not enclosed by any of %s' %
(node, loops_to_nodes_of_type))
self.builder.add_continue_node(node, try_node, guards)
def visit_ClassDef(self, node):
# We also keep the ClassDef node in the CFG, since it technically is a
# statement.
# For example, this is legal and allows executing user code:
#
# class Foo(bar()):
# pass
#
# It also has a scope:
#
# class Bar(object):
# a = 1
if self.builder is None:
self.generic_visit(node)
return
self.builder.add_ordinary_node(node)
self.builder_stack.append(self.builder)
self.builder = GraphBuilder(node)
self._enter_lexical_scope(node)
self._process_basic_statement(node)
self._exit_lexical_scope(node)
# TODO(mdan): Track the CFG local to the class definition as well?
self.builder = self.builder_stack.pop()
def _process_function_def(self, node, is_lambda):
# The function body is stored in a separate graph, because function
# definitions have effects very different from function calls.
if self.builder is not None:
self.builder.add_ordinary_node(node)
self.builder_stack.append(self.builder)
self.builder = GraphBuilder(node)
self._enter_lexical_scope(node)
self.builder.enter_section(node)
self._process_basic_statement(node.args)
if is_lambda:
self._process_exit_statement(node.body, (gast.Lambda,))
else:
for stmt in node.body:
self.visit(stmt)
self.builder.exit_section(node)
self._exit_lexical_scope(node)
self.cfgs[node] = self.builder.build()
self.builder = self.builder_stack.pop()
def visit_FunctionDef(self, node):
self._process_function_def(node, is_lambda=False)
def visit_Lambda(self, node):
self._process_function_def(node, is_lambda=True)
def visit_Return(self, node):
self._process_exit_statement(node, (gast.FunctionDef,))
def visit_Import(self, node):
self._process_basic_statement(node)
def visit_ImportFrom(self, node):
self._process_basic_statement(node)
def visit_Expr(self, node):
self._process_basic_statement(node)
def visit_NamedExpr(self, node):
# TODO(yileiyang): Add a test case once we have a newer astunparse version.
# NamedExpr was introduced in Python 3.8 and supported in gast 0.5.1+.
self._process_basic_statement(node)
def visit_Assign(self, node):
self._process_basic_statement(node)
def visit_AnnAssign(self, node):
self._process_basic_statement(node)
def visit_AugAssign(self, node):
self._process_basic_statement(node)
def visit_Pass(self, node):
self._process_basic_statement(node)
def visit_Global(self, node):
self._process_basic_statement(node)
def visit_Nonlocal(self, node):
self._process_basic_statement(node)
def visit_Print(self, node):
self._process_basic_statement(node)
def visit_Raise(self, node):
self._process_exit_statement(
node, (gast.FunctionDef,), may_exit_via_except=True)
self.builder.errors.add(node)
def visit_Assert(self, node):
# Ignoring the effect of exceptions.
self._process_basic_statement(node)
def visit_Delete(self, node):
self._process_basic_statement(node)
def visit_If(self, node):
# No need to track ifs as lexical scopes, for now.
# Lexical scopes are generally tracked in order to be able to resolve the
# targets of jump statements like break/continue/etc. Since there is no
# statement that can interrupt a conditional, we don't need to track their
# lexical scope. That may change in the future.
self.builder.begin_statement(node)
self.builder.enter_cond_section(node)
self._process_basic_statement(node.test)
self.builder.new_cond_branch(node)
for stmt in node.body:
self.visit(stmt)
self.builder.new_cond_branch(node)
for stmt in node.orelse:
self.visit(stmt)
self.builder.exit_cond_section(node)
self.builder.end_statement(node)
def visit_While(self, node):
self.builder.begin_statement(node)
self._enter_lexical_scope(node)
self.builder.enter_section(node)
self.generic_visit(node.test)
self.builder.enter_loop_section(node, node.test)
for stmt in node.body:
self.visit(stmt)
self.builder.exit_loop_section(node)
# Note: although the orelse is technically part of the loop node,
# the statements inside it don't affect the loop itself. For example, a
# break in the loop's orelse will not affect the loop itself.
self._exit_lexical_scope(node)
for stmt in node.orelse:
self.visit(stmt)
self.builder.exit_section(node)
self.builder.end_statement(node)
def visit_For(self, node):
self.builder.begin_statement(node)
self._enter_lexical_scope(node)
self.builder.enter_section(node)
# Note: Strictly speaking, this should be node.target + node.iter.
# However, the activity analysis accounts for this inconsistency,
# so dataflow analysis produces the correct values.
self.generic_visit(node.iter)
self.builder.enter_loop_section(node, node.iter)
# Also include the "extra loop test" annotation, to capture things like the
# control variable for return and break in for loops.
if anno.hasanno(node, anno.Basic.EXTRA_LOOP_TEST):
self._process_basic_statement(
anno.getanno(node, anno.Basic.EXTRA_LOOP_TEST))
for stmt in node.body:
self.visit(stmt)
self.builder.exit_loop_section(node)
# Note: although the orelse is technically part of the loop node,
# they don't count as loop bodies. For example, a break in the loop's
# orelse will affect the parent loop, not the current one.
self._exit_lexical_scope(node)
for stmt in node.orelse:
self.visit(stmt)
self.builder.exit_section(node)
self.builder.end_statement(node)
def visit_Break(self, node):
self._process_exit_statement(node, (
gast.While,
gast.For,
))
def visit_Continue(self, node):
self._process_continue_statement(node, (
gast.While,
gast.For,
))
def visit_ExceptHandler(self, node):
self.builder.begin_statement(node)
self.builder.enter_except_section(node)
if node.type is not None:
self.visit(node.type)
if node.name is not None:
self.visit(node.name)
for stmt in node.body:
self.visit(stmt)
self.builder.end_statement(node)
def visit_Try(self, node):
self.builder.begin_statement(node)
self._enter_lexical_scope(node)
# Note: the current simplification is that the try block fully executes
# regardless of whether an exception triggers or not. This is consistent
# with blocks free of try/except, which also don't account for the
# possibility of an exception being raised mid-block.
for stmt in node.body:
self.visit(stmt)
# The orelse is an optional continuation of the body.
if node.orelse:
block_representative = node.orelse[0]
self.builder.enter_cond_section(block_representative)
self.builder.new_cond_branch(block_representative)
for stmt in node.orelse:
self.visit(stmt)
self.builder.new_cond_branch(block_representative)
self.builder.exit_cond_section(block_representative)
self._exit_lexical_scope(node)
if node.handlers:
# Using node would be inconsistent. Using the first handler node is also
# inconsistent, but less so.
block_representative = node.handlers[0]
self.builder.enter_cond_section(block_representative)
for block in node.handlers:
self.builder.new_cond_branch(block_representative)
self.visit(block)
self.builder.new_cond_branch(block_representative)
self.builder.exit_cond_section(block_representative)
if node.finalbody:
self.builder.enter_finally_section(node)
for stmt in node.finalbody:
self.visit(stmt)
self.builder.exit_finally_section(node)
self.builder.end_statement(node)
def visit_With(self, node):
# TODO(mdan): Mark the context manager's exit call as exit guard.
for item in node.items:
self._process_basic_statement(item)
for stmt in node.body:
self.visit(stmt)
def build(node):
visitor = AstToCfg()
visitor.visit(node)
return visitor.cfgs
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,38 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
py_library(
name = "common_transformers",
srcs = [
"anf.py",
],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/python/autograph/pyct:gast_util",
"//tensorflow/python/autograph/pyct:templates",
"//tensorflow/python/autograph/pyct:transformer",
"@pypi//gast",
],
)
py_test(
name = "anf_test",
srcs = ["anf_test.py"],
strict_deps = True,
tags = ["no_oss"],
deps = [
":common_transformers",
"@pypi//gast",
#internal proto upb dep
"//tensorflow/python/autograph/pyct:loader",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:transformer",
"//tensorflow/python/platform:client_testlib",
],
)
@@ -0,0 +1,620 @@
# Copyright 2017 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.
# ==============================================================================
"""Conversion to A-normal form.
The general idea of A-normal form is that every intermediate value is
explicitly named with a variable. For more, see
https://en.wikipedia.org/wiki/A-normal_form.
The specific converters used here are based on Python AST semantics as
documented at https://greentreesnakes.readthedocs.io/en/latest/.
"""
import collections
import gast
from tensorflow.python.autograph.pyct import gast_util
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.autograph.pyct import transformer
# TODO(mdan): Replace with naming.Namer.
class DummyGensym:
"""A dumb gensym that suffixes a stem by sequential numbers from 1000."""
def __init__(self):
# A proper implementation needs to account for:
# * ctx.info.namespace
# * all the symbols defined in the AST
# * the symbols generated so far
self._idx = 0
def new_name(self, stem='tmp'):
self._idx += 1
return stem + '_' + str(1000 + self._idx)
REPLACE = lambda _1, _2, _3: True
LEAVE = lambda _1, _2, _3: False
ANY = object()
class ASTEdgePattern(collections.namedtuple(
'ASTEdgePattern', ['parent', 'field', 'child'])):
"""A pattern defining a type of AST edge.
This consists of three components:
- The type of the parent node, checked with isinstance,
- The name of the field, checked with string equality, and
- The type of the child node, also checked with isinstance.
If all three match, the whole pattern is considered to match.
In all three slots, the special value `anf.ANY` is treated as "match
anything". The internal nodes are produced from the `gast` library rather
than the standard `ast` module, which may affect `isinstance` checks.
"""
__slots__ = ()
def matches(self, parent, field, child):
"""Computes whether this pattern matches the given edge."""
if self.parent is ANY or isinstance(parent, self.parent):
pass # OK
else:
return False
if self.field is ANY or field == self.field:
pass # OK
else:
return False
return self.child is ANY or isinstance(child, self.child)
class AnfTransformer(transformer.Base):
"""Performs the conversion to A-normal form (ANF)."""
# The algorithm is a postorder recursive tree walk. Any given node A may, in
# general, require creation of a series B of Assign statements, which compute
# and explicitly name the intermediate values needed to compute the value of
# A. If A was already a statement, it can be replaced with the sequence B +
# [A]. If A was an expression, B needs to be propagated up the tree until a
# statement is encountered. Since the `ast.NodeTransformer` framework makes
# no provision for subtraversals returning side information, this class
# accumulates the sequence B in an instance variable.
# The only other subtlety is that some Python statements (like `if`) have both
# expression fields (`test`) and statement list fields (`body` and `orelse`).
# Any additional assignments needed to name all the intermediate values in the
# `test` can be prepended to the `if` node, but assignments produced by
# processing the `body` and the `orelse` need to be kept together with them,
# and not accidentally lifted out of the `if`.
def __init__(self, ctx, config):
"""Creates an ANF transformer.
Args:
ctx: transformer.Context
config: Configuration
"""
super(AnfTransformer, self).__init__(ctx)
if config is None:
# These could be pulled out, but are generally considered to already be in
# A-normal form. Thus they are left in by default, but could be pulled
# out if the configuration calls for it.
if gast_util.GAST2:
literal_node_types = (
gast.Num, gast.Str, gast.Bytes, gast.NameConstant,
gast.Name # Name is here to cover True, False, and None in Python 2
)
elif gast_util.GAST3:
literal_node_types = (
gast.Constant,
gast.Name # Name is here to cover True, False, and None in Python 2
)
else:
assert False
self._overrides = [
(ASTEdgePattern(ANY, ANY, literal_node_types), LEAVE),
(ASTEdgePattern(ANY, ANY, gast.expr), REPLACE)]
else:
self._overrides = config
self._gensym = DummyGensym()
self._pending_statements = []
def _consume_pending_statements(self):
ans = self._pending_statements
self._pending_statements = []
return ans
def _add_pending_statement(self, stmt):
self._pending_statements.append(stmt)
def _match(self, pattern, parent, field, child):
if pattern is ANY:
return True
else:
return pattern.matches(parent, field, child)
def _should_transform(self, parent, field, child):
for pat, result in self._overrides:
if self._match(pat, parent, field, child):
return result(parent, field, child)
# Fell off the end of the pattern list: do not transform
return False
def _do_transform_node(self, node):
temp_name = self._gensym.new_name()
temp_assign = templates.replace(
'temp_name = expr', temp_name=temp_name, expr=node)[0]
self._add_pending_statement(temp_assign)
answer = templates.replace('temp_name', temp_name=temp_name)[0]
return answer
def _ensure_node_in_anf(self, parent, field, node):
"""Puts `node` in A-normal form, by replacing it with a variable if needed.
The exact definition of A-normal form is given by the configuration. The
parent and the incoming field name are only needed because the configuration
may be context-dependent.
Args:
parent: An AST node, the parent of `node`.
field: The field name under which `node` is the child of `parent`.
node: An AST node, potentially to be replaced with a variable reference.
Returns:
node: An AST node; the argument if transformation was not necessary,
or the new variable reference if it was.
"""
if node is None:
return node
if _is_trivial(node):
return node
if isinstance(node, list):
# If something's field was actually a list, e.g., variadic arguments.
return [self._ensure_node_in_anf(parent, field, n) for n in node]
if isinstance(node, gast.keyword):
node.value = self._ensure_node_in_anf(parent, field, node.value)
return node
if isinstance(node, (gast.Starred, gast.withitem, gast.slice)):
# These nodes aren't really extractable in their own right, but their
# subnodes might be. Propagate the parent and field name to the child
# nodes, instead of querying the configuration for children of, e.g.,
# gast.Starred.
return self._ensure_fields_in_anf(node, parent, field)
if self._should_transform(parent, field, node):
return self._do_transform_node(node)
else:
return node
def _ensure_fields_in_anf(self, node, parent=None, super_field=None):
for field in node._fields:
if field.startswith('__'):
continue
parent_supplied = node if parent is None else parent
field_supplied = field if super_field is None else super_field
setattr(node, field, self._ensure_node_in_anf(
parent_supplied, field_supplied, getattr(node, field)))
return node
def _visit_strict_statement(self, node, children_ok_to_transform=True):
assert not self._pending_statements
node = self.generic_visit(node)
if children_ok_to_transform:
self._ensure_fields_in_anf(node)
results = self._consume_pending_statements()
results.append(node)
return results
def _visit_trivial_only_statement(self, node, msg):
assert not self._pending_statements
node = self.generic_visit(node)
self._ensure_fields_in_anf(node)
if self._pending_statements:
raise ValueError(msg)
else:
return node
def _visit_strict_expression(self, node):
node = self.generic_visit(node)
self._ensure_fields_in_anf(node)
return node
def _visit_trivial_only_expression(self, node, msg):
k = len(self._pending_statements)
node = self.generic_visit(node)
self._ensure_fields_in_anf(node)
# This check relies on there being no opportunities to consume pending
# statements while traversing children of an expression.
if len(self._pending_statements) != k:
raise ValueError(msg)
else:
return node
# Note on code order: These are listed in the same order as the grammar
# elements on https://github.com/serge-sans-paille/gast
# FunctionDef, AsyncFunctionDef, and ClassDef should be correct by default.
def visit_Return(self, node):
return self._visit_strict_statement(node)
def visit_Delete(self, node):
return self._visit_strict_statement(node, children_ok_to_transform=False)
def visit_Assign(self, node):
return self._visit_strict_statement(node, children_ok_to_transform=False)
def visit_AugAssign(self, node):
return self._visit_strict_statement(node, children_ok_to_transform=False)
def visit_Print(self, node):
return self._visit_strict_statement(node)
def visit_For(self, node):
assert not self._pending_statements
# It's important to visit node.iter first, because any statements created
# thereby need to live outside the body.
self.visit(node.iter)
node.iter = self._ensure_node_in_anf(node, 'iter', node.iter)
iter_stmts = self._consume_pending_statements()
# This generic_visit will revisit node.iter, but that is correct because by
# this point the node.iter link has been checked. It may be somewhat
# expensive if the configuration didn't call for transforming node.iter, as
# then it may be large and will be uselessly transformed again. This
# behavior is what causes the documented effect that configuration callables
# may be invoked more than once of the same links; if the code is rewritten
# not to do that (anywhere), the docstring of `transform` should be updated.
node = self.generic_visit(node)
assert not self._pending_statements
iter_stmts.append(node)
return iter_stmts
def visit_AsyncFor(self, node):
msg = ('Nontrivial AsyncFor nodes not supported yet '
'(need to think through the semantics).')
return self._visit_trivial_only_statement(node, msg)
def visit_While(self, node):
assert not self._pending_statements
self.visit(node.test)
node.test = self._ensure_node_in_anf(node, 'test', node.test)
if self._pending_statements:
msg = ('While with nontrivial test not supported yet '
'(need to avoid precomputing the test).')
raise ValueError(msg)
# If traversing node.test yielded no statements extracted, the generic visit
# will do the right thing.
return self.generic_visit(node)
def visit_If(self, node):
assert not self._pending_statements
# It's important to visit node.test first, because any statements created
# thereby need to live outside the body.
self.visit(node.test)
node.test = self._ensure_node_in_anf(node, 'test', node.test)
condition_stmts = self._consume_pending_statements()
# This generic_visit will revisit node.test, but that is correct because by
# this point the node.test link has been checked. It may be somewhat
# expensive if the configuration didn't call for transforming node.test, as
# then it may be large and will be uselessly transformed again. This
# happens in several places.
node = self.generic_visit(node)
assert not self._pending_statements
condition_stmts.append(node)
return condition_stmts
def visit_With(self, node):
assert not self._pending_statements
# It's important to visit node.items first, because any statements created
# thereby need to live outside the body.
for item in node.items:
self.visit(item)
node.items = [self._ensure_node_in_anf(node, 'items', n)
for n in node.items]
contexts_stmts = self._consume_pending_statements()
# This generic_visit will revisit node.items, but that is correct because by
# this point the node.items link has been checked. It may be somewhat
# expensive if the configuration didn't call for transforming node.items, as
# then it may be large and will be uselessly transformed again. This
# happens in several places.
node = self.generic_visit(node)
assert not self._pending_statements
contexts_stmts.append(node)
return contexts_stmts
def visit_AsyncWith(self, node):
msg = ('Nontrivial AsyncWith nodes not supported yet '
'(need to think through the semantics).')
return self._visit_trivial_only_statement(node, msg)
def visit_Raise(self, node):
return self._visit_strict_statement(node)
# Try should be correct by default.
def visit_Assert(self, node):
msg = ('Nontrivial Assert nodes not supported yet '
'(need to avoid computing the test when assertions are off, and '
'avoid computing the irritant when the assertion does not fire).')
return self._visit_trivial_only_statement(node, msg)
# Import and ImportFrom should be correct by default.
def visit_Exec(self, node):
return self._visit_strict_statement(node)
# Global and Nonlocal should be correct by default.
def visit_Expr(self, node):
return self._visit_strict_statement(node, children_ok_to_transform=False)
# Pass, Break, and Continue should be correct by default.
def visit_BoolOp(self, node):
msg = ('Nontrivial BoolOp nodes not supported yet '
'(need to preserve short-circuiting semantics).')
return self._visit_trivial_only_expression(node, msg)
def visit_BinOp(self, node):
return self._visit_strict_expression(node)
def visit_UnaryOp(self, node):
return self._visit_strict_expression(node)
def visit_Lambda(self, node):
msg = ('Nontrivial Lambda nodes not supported '
'(cannot insert statements into lambda bodies).')
return self._visit_trivial_only_expression(node, msg)
def visit_IfExp(self, node):
msg = ('Nontrivial IfExp nodes not supported yet '
'(need to convert to If statement, to evaluate branches lazily '
'and insert statements into them).')
return self._visit_trivial_only_expression(node, msg)
def visit_Dict(self, node):
return self._visit_strict_expression(node)
def visit_Set(self, node):
return self._visit_strict_expression(node)
def visit_ListComp(self, node):
msg = ('ListComp nodes not supported '
'(need to convert to a form that tolerates '
'assignment statements in clause bodies).')
raise ValueError(msg)
def visit_SetComp(self, node):
msg = ('SetComp nodes not supported '
'(need to convert to a form that tolerates '
'assignment statements in clause bodies).')
raise ValueError(msg)
def visit_DictComp(self, node):
msg = ('DictComp nodes not supported '
'(need to convert to a form that tolerates '
'assignment statements in clause bodies).')
raise ValueError(msg)
def visit_GeneratorExp(self, node):
msg = ('GeneratorExp nodes not supported '
'(need to convert to a form that tolerates '
'assignment statements in clause bodies).')
raise ValueError(msg)
def visit_Await(self, node):
msg = ('Nontrivial Await nodes not supported yet '
'(need to think through the semantics).')
return self._visit_trivial_only_expression(node, msg)
def visit_Yield(self, node):
return self._visit_strict_expression(node)
def visit_YieldFrom(self, node):
msg = ('Nontrivial YieldFrom nodes not supported yet '
'(need to unit-test them in Python 2).')
return self._visit_trivial_only_expression(node, msg)
def visit_Compare(self, node):
if len(node.ops) > 1:
msg = ('Multi-ary compare nodes not supported yet '
'(need to preserve short-circuiting semantics).')
raise ValueError(msg)
return self._visit_strict_expression(node)
def visit_Call(self, node):
return self._visit_strict_expression(node)
def visit_Repr(self, node):
msg = ('Nontrivial Repr nodes not supported yet '
'(need to research their syntax and semantics).')
return self._visit_trivial_only_expression(node, msg)
def visit_FormattedValue(self, node):
msg = ('Nontrivial FormattedValue nodes not supported yet '
'(need to unit-test them in Python 2).')
return self._visit_trivial_only_expression(node, msg)
def visit_JoinedStr(self, node):
msg = ('Nontrivial JoinedStr nodes not supported yet '
'(need to unit-test them in Python 2).')
return self._visit_trivial_only_expression(node, msg)
def visit_Attribute(self, node):
return self._visit_strict_expression(node)
def visit_Subscript(self, node):
return self._visit_strict_expression(node)
# Starred and Name are correct by default, because the right thing to do is to
# just recur.
def visit_List(self, node):
node = self.generic_visit(node)
if not isinstance(node.ctx, gast.Store):
self._ensure_fields_in_anf(node)
return node
def visit_Tuple(self, node):
node = self.generic_visit(node)
if not isinstance(node.ctx, gast.Store):
self._ensure_fields_in_anf(node)
return node
def _is_py2_name_constant(node):
return isinstance(node, gast.Name) and node.id in ['True', 'False', 'None']
def _is_trivial(node):
"""Returns whether to consider the given node 'trivial'.
The definition of 'trivial' is a node that can't meaningfully be pulled out
into its own assignment statement.
This is surprisingly difficult to do robustly across versions of Python and
gast, as the parsing of constants has changed, if I may, constantly.
Args:
node: An AST node to check for triviality
Returns:
trivial: A Python `bool` indicating whether the node is trivial.
"""
trivial_node_types = (
# Variable names
gast.Name,
# Non-nodes that show up as AST fields
bool,
str,
# Binary operators
gast.Add,
gast.Sub,
gast.Mult,
gast.Div,
gast.Mod,
gast.Pow,
gast.LShift,
gast.RShift,
gast.BitOr,
gast.BitXor,
gast.BitAnd,
gast.FloorDiv,
# Unary operators
gast.Invert,
gast.Not,
gast.UAdd,
gast.USub,
# Comparison operators
gast.Eq,
gast.NotEq,
gast.Lt,
gast.LtE,
gast.Gt,
gast.GtE,
gast.Is,
gast.IsNot,
gast.In,
gast.NotIn,
# Other leaf nodes that don't make sense standalone.
gast.expr_context,
)
if isinstance(node, trivial_node_types) and not _is_py2_name_constant(node):
return True
if gast_util.is_ellipsis(node):
return True
return False
def transform(node, ctx, config=None):
"""Converts the given node to A-normal form (ANF).
The general idea of A-normal form: https://en.wikipedia.org/wiki/A-normal_form
The specific converters used here are based on Python AST semantics as
documented at https://greentreesnakes.readthedocs.io/en/latest/.
What exactly should be considered A-normal form for any given programming
language is not completely obvious. The transformation defined here is
therefore configurable as to which syntax to replace with a fresh variable and
which to leave be. The configuration is intentionally flexible enough to
define very precise variable insertion transformations, should that be
desired.
The configuration is a list of syntax rules, each of which is a 2-tuple:
- An `ASTEdgePattern` (which see) defining a type of AST edge, and
- Whether to transform children of such edges.
The special object `anf.ANY` may be used as a pattern that matches all edges.
Each replacement directive is one of three possible things:
- The object `anf.REPLACE`, meaning "Replace this child node with a variable",
- The object `anf.LEAVE`, meaning "Do not replace this child node with a
variable", or
- A Python callable. If a callable, it is called with the parent node, the
field name, and the child node, and must compute a boolean indicating
whether to transform the child node or not. The callable is free to use
whatever context information it chooses. The callable may be invoked more
than once on the same link, and must produce the same answer each time.
The syntax rules are tested in order, and the first match governs. If no rule
matches, the node is not transformed.
The above rules notwithstanding,
- Variable references are never replaced with (fresh) variables, as that would
accomplish nothing.
- The left-hand children of Assign and AugAssign nodes, and the children of
Del nodes, are never replaced with variables, as that would break their
semantics.
- The right-hand children of Assign nodes are never replaced with variables,
as the original assignment would still have to be present in the result
to define the new variable. (That is, there's no point in transforming
`x = sin(y)` into `tmp = sin(y); x = tmp`.)
- The right-hand children of AugAssign nodes are never replaced with variables
either, but only because the difference from Assign was considered a
potential source of confusion (and it would have been slightly awkward in
the code to treat the RHS differently than the LHS).
- Various special-purpose AST nodes are not exposed to the configuration, lest
the transform produce invalid syntax like, e.g., `tmp = +; x = 1 tmp 2`.
For example, the configuration
```python
[(anf.ASTEdgePattern(anf.ANY, anf.ANY, gast.expr), anf.REPLACE)]
```
gives explicit fresh names to all expressions regardless of context (except as
outlined above), whereas
```python
[(anf.ASTEdgePattern(gast.If, "test", anf.ANY), anf.REPLACE)]
```
only transforms the conditionals of `if` statements (but not, e.g., `while`).
If no configuration is supplied, the default behavior is to transform all
expressions except literal constants, which is defined as a configuration as
```python
# For Python 3, and gast library versions before 0.3
literals = (gast.Num, gast.Str, gast.Bytes, gast.NameConstant)
[(anf.ASTEdgePattern(anf.ANY, anf.ANY, literals), anf.LEAVE),
(anf.ASTEdgePattern(anf.ANY, anf.ANY, gast.expr), anf.REPLACE)]
```
Args:
node: The node to transform.
ctx: transformer.EntityInfo. TODO(mdan): What information does this
argument provide?
config: Optional ANF configuration. If omitted, ANF replaces all expression
expect literal constants.
"""
return AnfTransformer(ctx, config).visit(node)
@@ -0,0 +1,518 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for anf module."""
import textwrap
import gast
from tensorflow.python.autograph.pyct import loader
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import transformer
from tensorflow.python.autograph.pyct.common_transformers import anf
from tensorflow.python.platform import test
# TODO(mdan): These two functions no longer need to be at the top level.
# TODO(mdan): Don't use exec.
def exec_test_function():
# The point is to test A-normal form conversion of exec
# pylint: disable=exec-used
exec('computed' + 5 + 'stuff', globals(), locals())
def exec_expected_result():
# pylint: disable=exec-used
tmp_1001 = 'computed' + 5
tmp_1002 = tmp_1001 + 'stuff'
tmp_1003 = globals()
tmp_1004 = locals()
exec(tmp_1002, tmp_1003, tmp_1004)
class AnfTestBase(test.TestCase):
def _simple_context(self):
entity_info = transformer.EntityInfo(
name='test_fn',
source_code=None,
source_file=None,
future_features=(),
namespace=None)
return transformer.Context(entity_info, None, None)
def assert_same_ast(self, expected_node, node, msg=None):
expected_source = parser.unparse(expected_node, indentation=' ')
expected_str = textwrap.dedent(expected_source).strip()
got_source = parser.unparse(node, indentation=' ')
got_str = textwrap.dedent(got_source).strip()
self.assertEqual(expected_str, got_str, msg=msg)
def assert_body_anfs_as_expected(self, expected_fn, test_fn, config=None):
# Testing the code bodies only. Wrapping them in functions so the
# syntax highlights nicely, but Python doesn't try to execute the
# statements.
exp_node, _ = parser.parse_entity(expected_fn, future_features=())
node, _ = parser.parse_entity(test_fn, future_features=())
node = anf.transform(node, self._simple_context(), config=config)
exp_name = exp_node.name
# Ignoring the function names in the result because they can't be
# the same (because both functions have to exist in the same scope
# at the same time).
node.name = exp_name
self.assert_same_ast(exp_node, node)
# Check that ANF is idempotent
node_repeated = anf.transform(node, self._simple_context())
self.assert_same_ast(node_repeated, node)
class AnfTransformerTest(AnfTestBase):
def test_basic(self):
def test_function():
a = 0
return a
node, _ = parser.parse_entity(test_function, future_features=())
node = anf.transform(node, self._simple_context())
result, _, _ = loader.load_ast(node)
self.assertEqual(test_function(), result.test_function())
def test_binop_basic(self):
def test_function(x, y, z):
a = x + y + z
return a
def expected_result(x, y, z):
tmp_1001 = x + y
a = tmp_1001 + z
return a
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_if_basic(self):
def test_function(a, b, c, e, f, g):
if a + b + c:
d = e + f + g
return d
def expected_result(a, b, c, e, f, g):
tmp_1001 = a + b
tmp_1002 = tmp_1001 + c
if tmp_1002:
tmp_1003 = e + f
d = tmp_1003 + g
return d
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_nested_binop_and_return(self):
def test_function(b, c, d, e):
return (2 * b + c) + (d + e)
def expected_result(b, c, d, e):
tmp_1001 = 2 * b
tmp_1002 = tmp_1001 + c
tmp_1003 = d + e
tmp_1004 = tmp_1002 + tmp_1003
return tmp_1004
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_function_call_and_expr(self):
def test_function(call_something, a, b, y, z, c, d, e, f, g, h, i):
call_something(a + b, y * z, kwarg=c + d, *(e + f), **(g + h + i))
def expected_result(call_something, a, b, y, z, c, d, e, f, g, h, i):
tmp_1001 = g + h
tmp_1002 = a + b
tmp_1003 = y * z
tmp_1004 = e + f
tmp_1005 = c + d
tmp_1006 = tmp_1001 + i
call_something(tmp_1002, tmp_1003, kwarg=tmp_1005, *tmp_1004, **tmp_1006)
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_with_and_print(self):
def test_function(a, b, c):
with a + b + c as d:
print(2 * d + 1)
def expected_result(a, b, c):
tmp_1001 = a + b
tmp_1002 = tmp_1001 + c
with tmp_1002 as d:
tmp_1003 = 2 * d
tmp_1004 = tmp_1003 + 1
print(tmp_1004)
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_nested_multi_value_assign(self):
def test_function(a, b, c):
x, y = a, a + b
(z, y), x = (c, y + b), x + a
return z, (y, x)
def expected_result(a, b, c):
tmp_1001 = a + b
x, y = a, tmp_1001
tmp_1002 = y + b
tmp_1003 = (c, tmp_1002)
tmp_1004 = x + a
(z, y), x = tmp_1003, tmp_1004
tmp_1005 = y, x
tmp_1006 = z, tmp_1005
return tmp_1006
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_deeply_nested_multi_value_assign(self):
def test_function(a):
[([(b, c), [d, e]], (f, g)), [(h, i, j), k]] = a
return [([(b, c), [d, e]], (f, g)), [(h, i, j), k]]
def expected_result(a):
[([(b, c), [d, e]], (f, g)), [(h, i, j), k]] = a
tmp_1001 = b, c
tmp_1002 = [d, e]
tmp_1003 = [tmp_1001, tmp_1002]
tmp_1004 = f, g
tmp_1005 = h, i, j
tmp_1006 = tmp_1003, tmp_1004
tmp_1007 = [tmp_1005, k]
tmp_1008 = [tmp_1006, tmp_1007]
return tmp_1008
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_local_definition_and_binary_compare(self):
def test_function():
def foo(a, b):
return 2 * a < b
return foo
def expected_result():
def foo(a, b):
tmp_1001 = 2 * a
tmp_1002 = tmp_1001 < b
return tmp_1002
return foo
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_list_literal(self):
def test_function(a, b, c, d, e, f):
return [a + b, c + d, e + f]
def expected_result(a, b, c, d, e, f):
tmp_1001 = a + b
tmp_1002 = c + d
tmp_1003 = e + f
tmp_1004 = [tmp_1001, tmp_1002, tmp_1003]
return tmp_1004
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_tuple_literal_and_unary(self):
def test_function(a, b, c, d, e, f):
return (a + b, -(c + d), e + f)
def expected_result(a, b, c, d, e, f):
tmp_1001 = c + d
tmp_1002 = a + b
tmp_1003 = -tmp_1001
tmp_1004 = e + f
tmp_1005 = (tmp_1002, tmp_1003, tmp_1004)
return tmp_1005
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_set_literal(self):
def test_function(a, b, c, d, e, f):
return set(a + b, c + d, e + f)
def expected_result(a, b, c, d, e, f):
tmp_1001 = a + b
tmp_1002 = c + d
tmp_1003 = e + f
tmp_1004 = set(tmp_1001, tmp_1002, tmp_1003)
return tmp_1004
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_dict_literal_and_repr(self):
def test_function(foo, bar, baz):
return repr({foo + bar + baz: 7 | 8})
def expected_result(foo, bar, baz):
tmp_1001 = foo + bar
tmp_1002 = tmp_1001 + baz
tmp_1003 = 7 | 8
tmp_1004 = {tmp_1002: tmp_1003}
tmp_1005 = repr(tmp_1004)
return tmp_1005
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_field_read_and_write(self):
def test_function(a, d):
a.b.c = d.e.f + 3
def expected_result(a, d):
tmp_1001 = a.b
tmp_1002 = d.e
tmp_1003 = tmp_1002.f
tmp_1001.c = tmp_1003 + 3
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_subscript_read_and_write(self):
def test_function(a, b, c, d, e, f):
a[b][c] = d[e][f] + 3
def expected_result(a, b, c, d, e, f):
tmp_1001 = a[b]
tmp_1002 = d[e]
tmp_1003 = tmp_1002[f]
tmp_1001[c] = tmp_1003 + 3
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_augassign_and_delete(self):
def test_function(a, x, y, z):
a += x + y + z
del a
del z[y][x]
def expected_result(a, x, y, z):
tmp_1001 = x + y
a += tmp_1001 + z
del a
tmp_1002 = z[y]
del tmp_1002[x]
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_raise_yield_and_raise(self):
def test_function(a, c, some_computed, exception):
yield a ** c
raise some_computed('complicated' + exception)
def expected_result(a, c, some_computed, exception):
tmp_1001 = a ** c
yield tmp_1001
tmp_1002 = 'complicated' + exception
tmp_1003 = some_computed(tmp_1002)
raise tmp_1003
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_with_and_if_with_expressions(self):
def test_function(foo, bar, function, quux, quozzle, w, x, y, z):
with foo + bar:
function(x + y)
if quux + quozzle:
function(z / w)
def expected_result(foo, bar, function, quux, quozzle, w, x, y, z):
tmp_1001 = foo + bar
with tmp_1001:
tmp_1002 = x + y
function(tmp_1002)
tmp_1003 = quux + quozzle
if tmp_1003:
tmp_1004 = z / w
function(tmp_1004)
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_exec(self):
self.assert_body_anfs_as_expected(exec_expected_result, exec_test_function)
def test_simple_while_and_assert(self):
def test_function(foo, quux):
while foo:
assert quux
foo = foo + 1 * 3
def expected_result(foo, quux):
while foo:
assert quux
tmp_1001 = 1 * 3
foo = foo + tmp_1001
self.assert_body_anfs_as_expected(expected_result, test_function)
def test_for(self):
def test_function(compute, something, complicated, foo):
for foo in compute(something + complicated):
bar = foo + 1 * 3
return bar
def expected_result(compute, something, complicated, foo):
tmp_1001 = something + complicated
tmp_1002 = compute(tmp_1001)
for foo in tmp_1002:
tmp_1003 = 1 * 3
bar = foo + tmp_1003
return bar
self.assert_body_anfs_as_expected(expected_result, test_function)
# This test collects several examples where the definition of A-normal form
# implemented by this transformer is questionable. Mostly it's here to spell
# out what the definition is in these cases.
def test_controversial(self):
def test_function(b, c, d, f):
a = c + d
a.b = c + d
a[b] = c + d
a += c + d
a, b = c
a, b = c, d
a = f(c)
a = f(c + d)
a[b + d] = f.e(c + d)
def expected_result(b, c, d, f):
a = c + d
a.b = c + d # Should be a.b = tmp? (Definitely not tmp = c + d)
a[b] = c + d # Should be a[b] = tmp? (Definitely not tmp = c + d)
a += c + d # Should be a += tmp? (Definitely not tmp = c + d)
a, b = c # Should be a = c[0], b = c[1]? Or not?
a, b = c, d # Should be a = c, b = d? Or not?
a = f(c)
tmp_1001 = c + d
a = f(tmp_1001)
tmp_1002 = b + d
tmp_1003 = f.e
tmp_1004 = c + d
a[tmp_1002] = tmp_1003(tmp_1004) # Or should be a[tmp1] = tmp2?
self.assert_body_anfs_as_expected(expected_result, test_function)
class AnfNonTransformationTest(AnfTransformerTest):
"""Test that specifying "no transformation" does nothing.
Reuses all the examples of AnfTransformerTest by overriding
`assert_body_anfs_as_expected_`.
"""
def assert_body_anfs_as_expected(self, expected_fn, test_fn):
# Testing the code bodies only. Wrapping them in functions so the
# syntax highlights nicely, but Python doesn't try to execute the
# statements.
node, _ = parser.parse_entity(test_fn, future_features=())
orig_source = parser.unparse(node, indentation=' ')
orig_str = textwrap.dedent(orig_source).strip()
config = [(anf.ANY, anf.LEAVE)] # Configuration to transform nothing
node = anf.transform(node, self._simple_context(), config=config)
new_source = parser.unparse(node, indentation=' ')
new_str = textwrap.dedent(new_source).strip()
self.assertEqual(orig_str, new_str)
class AnfConfiguredTest(AnfTestBase):
def test_constants_in_function_calls(self):
# An example specific configuration that differs from the default: Moving
# literals out of being directly passed to functions, but nothing else.
try:
# TODO(b/140808434): Fix this.
# gast pre-0.3
literals = (gast.Num, gast.Str, gast.Bytes, gast.NameConstant, gast.Name)
except AttributeError:
# gast 0.3+
literals = (gast.Constant, gast.Name)
config = [(anf.ASTEdgePattern(gast.Call, anf.ANY, literals), anf.REPLACE)]
def test_function(x, frob):
return frob(x, x+1, 2)
def expected_result(x, frob):
tmp_1001 = 2
return frob(x, x+1, tmp_1001)
self.assert_body_anfs_as_expected(expected_result, test_function, config)
def test_anf_some_function_calls(self):
# Another example specific configuration that differs from the default:
# Moving all arguments out of some function calls but leaving others be.
allowlist = ['foo']
def transform(parent, field, child):
del field
del child
func_name = parent.func.id
return str(func_name) in allowlist
config = [(anf.ASTEdgePattern(gast.Call, anf.ANY, anf.ANY), transform)]
def test_function(x, foo, bar):
y = foo(x, x+1, 2)
return bar(y, y+1, 2)
def expected_result(x, foo, bar):
tmp_1001 = x+1
tmp_1002 = 2
y = foo(x, tmp_1001, tmp_1002)
return bar(y, y+1, 2)
self.assert_body_anfs_as_expected(expected_result, test_function, config)
def test_touching_name_constant(self):
# Checking that the nodes for `True`, `False`, and `None` can be manipulated
# by a configuration. This is non-trivial, because in Python 2 those are
# represented as `Name`, which is the same node type as variable references.
specials = (gast.Name, gast.Constant)
config = [(anf.ASTEdgePattern(gast.Call, anf.ANY, specials), anf.REPLACE)]
def test_function(f):
return f(True, False, None)
def expected_result(f):
tmp_1001 = True
tmp_1002 = False
tmp_1003 = None
return f(tmp_1001, tmp_1002, tmp_1003)
self.assert_body_anfs_as_expected(expected_result, test_function, config)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,230 @@
# Copyright 2017 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.
# ==============================================================================
"""Code transformation exceptions."""
import collections
from tensorflow.python.autograph.pyct import origin_info
from tensorflow.python.util import traceback_utils
class FrameInfo(
collections.namedtuple('FrameInfo',
('filename', 'lineno', 'function_name', 'code',
'is_converted', 'is_allowlisted'))):
__slots__ = ()
def _stack_trace_inside_mapped_code(tb, source_map, converter_filename):
"""Summarizes inner traceback frames up to the call to a given function.
This functions locates the innermost (i.e. most recent) frame that corresponds
to code that can be mapped by source_map originated from, and returns a
translated stack trace ending at that frame. If no such frame is found, the
entire stack trace is summarized.
For example, the following code:
def f():
for i in tf.range(1):
z = y + i # z only defined here
Would generate this traceback:
<converted code>
ag__.for_stmt(...)
<for_stmt>
return _known_len_tf_for_stmt(iter_, extra_test, body, init_state)
<_known_len_tf_for_stmt>
_disallow_undefs_into_loop(*init_state)
<_disallow_undefs_into_loop>
raise ...
Which is then processed into:
<f>
for i in tf.range(1):
<for_stmt>
return _known_len_tf_for_stmt(iter_, extra_test, body, init_state)
<_known_len_tf_for_stmt>
_disallow_undefs_into_loop(*init_state)
<_disallow_undefs_into_loop>
raise ...
Args:
tb: traceback.FrameSummary, The traceback corresponding to an error.
Typically, the output of traceback.Summary.extract(capture_locals=True).
source_map: Dict[LineLocation, OriginInfo], a source map as created by
origin_info.create_source_map.
converter_filename: str, the file path of the converted module. Call frames
corresponding to this module are elided and their preceding frames are
marked as allowlisted. Note that frames enclosing converted code are
dropped using a different mechanism.
Returns:
List[FrameInfo]
"""
result_frames = []
for filename, line_number, function_name, text in reversed(tb):
loc = origin_info.LineLocation(filename=filename, lineno=line_number)
if loc in source_map:
origin = source_map[loc]
fi = FrameInfo(
filename=origin.loc.filename,
lineno=origin.loc.lineno,
function_name=origin.function_name,
code=origin.source_code_line,
is_converted=True,
is_allowlisted=False)
result_frames.append(fi)
break
if filename == converter_filename:
if result_frames:
prev = result_frames[-1]
assert not prev.is_converted # See the if above.
fi = FrameInfo(
filename=prev.filename,
lineno=prev.lineno,
function_name=prev.function_name,
code=prev.code,
is_converted=False,
is_allowlisted=True)
result_frames[-1] = fi
continue
fi = FrameInfo(
filename=filename,
lineno=line_number,
function_name=function_name,
code=text,
is_converted=False,
is_allowlisted=False)
result_frames.append(fi)
return tuple(result_frames)
KNOWN_STRING_CONSTRUCTOR_ERRORS = (
AssertionError,
AttributeError,
NameError,
NotImplementedError,
RuntimeError,
StopIteration,
TypeError,
UnboundLocalError,
ValueError,
)
# KeyError escapes newlines in strings. We create a special subclass
# that doesn't do that. Overriding the name for display purposes; hopefully
# that won't create too many surprises.
class MultilineMessageKeyError(KeyError):
def __init__(self, message, original_key):
super(MultilineMessageKeyError, self).__init__(original_key)
self.__message = message
def __str__(self):
return self.__message
MultilineMessageKeyError.__name__ = KeyError.__name__
class ErrorMetadataBase(object):
"""Container objects attached to exceptions raised in user code.
This metadata allows re-raising exceptions that occur in generated code, with
a custom error message that includes a stack trace relative to user-readable
code from which the generated code originated.
"""
__slots__ = ('translated_stack', 'cause_message')
def __init__(self, callsite_tb, cause_metadata, cause_message, source_map,
converter_filename):
translated_stack = _stack_trace_inside_mapped_code(
callsite_tb, source_map, converter_filename)
if cause_metadata is None:
self.translated_stack = translated_stack
self.cause_message = cause_message
else:
# Daisy chain the translated stacks.
self.translated_stack = (
cause_metadata.translated_stack + (translated_stack[-1],))
self.cause_message = cause_metadata.cause_message
def get_message(self):
"""Returns the message for the underlying exception."""
lines = []
lines.append('in user code:')
lines.append('')
for frame_info in reversed(self.translated_stack):
if (traceback_utils.is_traceback_filtering_enabled() and
not traceback_utils.include_frame(frame_info.filename)):
continue
# Same format with Python traceback.
formatted_line = (f' File "{frame_info.filename}", line '
f'{frame_info.lineno}, in {frame_info.function_name}')
if frame_info.is_converted:
formatted_line += ' *'
elif frame_info.is_allowlisted:
formatted_line += ' **'
lines.append(formatted_line)
if frame_info.code is None:
code_snippet = '<source unavailable>'
else:
code_snippet = frame_info.code.strip()
lines.append(' {}'.format(code_snippet))
lines.append('')
message_lines = self.cause_message.split('\n')
for i in range(len(message_lines)):
message_lines[i] = ' ' + message_lines[i]
lines.extend(message_lines)
lines.append('')
return '\n'.join(lines)
def create_exception(self, source_error):
"""Creates exception from source_error."""
preferred_type = type(source_error)
to_ret = None
if preferred_type.__init__ is Exception.__init__:
to_ret = preferred_type(self.get_message())
if preferred_type in KNOWN_STRING_CONSTRUCTOR_ERRORS:
to_ret = preferred_type(self.get_message())
elif preferred_type is KeyError:
to_ret = MultilineMessageKeyError(self.get_message(), self.cause_message)
if to_ret is not None:
return to_ret.with_traceback(source_error.__traceback__)
def to_exception(self, source_error):
exc = self.create_exception(source_error)
exc.__suppress_context__ = True
exc.ag_error_metadata = self
return exc
@@ -0,0 +1,126 @@
# Copyright 2019 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.
# ==============================================================================
"""Tests for error_utils module."""
import re
from tensorflow.python.autograph.pyct import error_utils
from tensorflow.python.autograph.pyct import origin_info
from tensorflow.python.platform import test
class ErrorMetadataBaseTest(test.TestCase):
def test_create_exception_default_constructor(self):
class CustomError(Exception):
pass
em = error_utils.ErrorMetadataBase(
callsite_tb=(),
cause_metadata=None,
cause_message='test message',
source_map={},
converter_filename=None)
exc = em.create_exception(CustomError())
self.assertIsInstance(exc, CustomError)
self.assertIn('test message', str(exc))
def test_create_exception_custom_constructor(self):
class CustomError(Exception):
def __init__(self):
super(CustomError, self).__init__('test_message')
em = error_utils.ErrorMetadataBase(
callsite_tb=(),
cause_metadata=None,
cause_message='test message',
source_map={},
converter_filename=None)
exc = em.create_exception(CustomError())
self.assertIsNone(exc)
def test_get_message_no_code(self):
callsite_tb = [
('/path/one.py', 11, 'test_fn_1', None),
('/path/two.py', 171, 'test_fn_2', 'test code'),
]
cause_message = 'Test message'
em = error_utils.ErrorMetadataBase(
callsite_tb=callsite_tb,
cause_metadata=None,
cause_message=cause_message,
source_map={},
converter_filename=None)
self.assertRegex(
em.get_message(),
re.compile(('"/path/one.py", line 11, in test_fn_1.*'
'"/path/two.py", line 171, in test_fn_2.*'
'Test message'), re.DOTALL))
def test_get_message_converted_code(self):
callsite_tb = [
('/path/one.py', 11, 'test_fn_1', 'test code 1'),
('/path/two.py', 171, 'test_fn_2', 'test code 2'),
('/path/three.py', 171, 'test_fn_3', 'test code 3'),
]
cause_message = 'Test message'
em = error_utils.ErrorMetadataBase(
callsite_tb=callsite_tb,
cause_metadata=None,
cause_message=cause_message,
source_map={
origin_info.LineLocation(filename='/path/two.py', lineno=171):
origin_info.OriginInfo(
loc=origin_info.LineLocation(
filename='/path/other_two.py', lineno=13),
function_name='converted_fn',
source_code_line='converted test code',
comment=None)
},
converter_filename=None)
result = em.get_message()
self.assertRegex(
result,
re.compile((r'converted_fn \*.*'
r'"/path/three.py", line 171, in test_fn_3.*'
r'Test message'), re.DOTALL))
self.assertNotRegex(result, re.compile('test_fn_1'))
def test_get_message_call_overload(self):
callsite_tb = [
('/path/one.py', 11, 'test_fn_1', 'test code 1'),
('/path/two.py', 0, 'test_fn_2', 'test code 2'),
('/path/three.py', 171, 'test_fn_3', 'test code 3'),
]
cause_message = 'Test message'
em = error_utils.ErrorMetadataBase(
callsite_tb=callsite_tb,
cause_metadata=None,
cause_message=cause_message,
source_map={},
converter_filename='/path/two.py')
self.assertRegex(
em.get_message(),
re.compile((r'"/path/one.py", line 11, in test_fn_1.*'
r'"/path/three.py", line 171, in test_fn_3 \*\*.*'
r'Test message'), re.DOTALL))
if __name__ == '__main__':
test.main()
@@ -0,0 +1,27 @@
# Copyright 2017 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.
# ==============================================================================
"""Code transformation exceptions."""
class PyCTError(Exception):
"""Base class for all exceptions."""
class UnsupportedLanguageElementError(PyCTError, NotImplementedError):
"""Raised for code patterns that AutoGraph does not support."""
class InaccessibleSourceCodeError(PyCTError, ValueError):
"""Raised when inspect can not access source code."""
@@ -0,0 +1,74 @@
# Copyright 2017 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.
# ==============================================================================
"""Gast compatibility library. Supports 0.2.2 and 0.3.2."""
# TODO(mdan): Remove this file once it's safe to break compatibility.
import functools
import gast
GAST2 = hasattr(gast, 'Str')
GAST3 = not GAST2
def _is_constant_gast_2(node):
return isinstance(node, (gast.Num, gast.Str, gast.Bytes, gast.Ellipsis,
gast.NameConstant))
def _is_constant_gast_3(node):
return isinstance(node, gast.Constant)
def is_literal(node):
"""Tests whether node represents a Python literal."""
# Normal literals, True/False/None/Etc. in Python3
if is_constant(node):
return True
# True/False/None/Etc. in Python2
if isinstance(node, gast.Name) and node.id in ['True', 'False', 'None']:
return True
return False
def _is_ellipsis_gast_2(node):
return isinstance(node, gast.Ellipsis)
def _is_ellipsis_gast_3(node):
return isinstance(node, gast.Constant) and node.value == Ellipsis
if GAST2:
is_constant = _is_constant_gast_2
is_ellipsis = _is_ellipsis_gast_2
Module = gast.Module
Name = gast.Name
Str = gast.Str
elif GAST3:
is_constant = _is_constant_gast_3
is_ellipsis = _is_ellipsis_gast_3
Module = functools.partial(gast.Module, type_ignores=None) # pylint:disable=invalid-name
Name = functools.partial(gast.Name, type_comment=None) # pylint:disable=invalid-name
Str = functools.partial(gast.Constant, kind=None) # pylint:disable=invalid-name
else:
assert False
@@ -0,0 +1,321 @@
# Copyright 2017 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.
# ==============================================================================
"""Live entity inspection utilities.
This module contains whatever inspect doesn't offer out of the box.
"""
import builtins
import inspect
import itertools
import linecache
import sys
import threading
import types
from tensorflow.python.util import tf_inspect
# This lock seems to help avoid linecache concurrency errors.
_linecache_lock = threading.Lock()
# Cache all the builtin elements in a frozen set for faster lookup.
_BUILTIN_FUNCTION_IDS = frozenset(id(v) for v in builtins.__dict__.values())
def islambda(f):
if not tf_inspect.isfunction(f):
return False
# TODO(mdan): Look into checking the only the code object.
if not (hasattr(f, '__name__') and hasattr(f, '__code__')):
return False
# Some wrappers can rename the function, but changing the name of the
# code object is harder.
return ((f.__name__ == '<lambda>') or (f.__code__.co_name == '<lambda>'))
def isnamedtuple(f):
"""Returns True if the argument is a namedtuple-like."""
if not (tf_inspect.isclass(f) and issubclass(f, tuple)):
return False
if not hasattr(f, '_fields'):
return False
fields = getattr(f, '_fields')
if not isinstance(fields, tuple):
return False
if not all(isinstance(f, str) for f in fields):
return False
return True
def isbuiltin(f):
"""Returns True if the argument is a built-in function."""
if id(f) in _BUILTIN_FUNCTION_IDS:
return True
elif isinstance(f, types.BuiltinFunctionType):
return True
elif inspect.isbuiltin(f):
return True
elif f is eval:
return True
else:
return False
def isconstructor(cls):
"""Returns True if the argument is an object constructor.
In general, any object of type class is a constructor, with the exception
of classes created using a callable metaclass.
See below for why a callable metaclass is not a trivial combination:
https://docs.python.org/2.7/reference/datamodel.html#customizing-class-creation
Args:
cls: Any
Returns:
Bool
"""
return (inspect.isclass(cls) and
not (issubclass(cls.__class__, type) and
hasattr(cls.__class__, '__call__') and
cls.__class__.__call__ is not type.__call__))
def _fix_linecache_record(obj):
"""Fixes potential corruption of linecache in the presence of functools.wraps.
functools.wraps modifies the target object's __module__ field, which seems
to confuse linecache in special instances, for example when the source is
loaded from a .par file (see https://google.github.io/subpar/subpar.html).
This function simply triggers a call to linecache.updatecache when a mismatch
was detected between the object's __module__ property and the object's source
file.
Args:
obj: Any
"""
if hasattr(obj, '__module__'):
obj_file = inspect.getfile(obj)
obj_module = obj.__module__
# A snapshot of the loaded modules helps avoid "dict changed size during
# iteration" errors.
loaded_modules = tuple(sys.modules.values())
for m in loaded_modules:
if hasattr(m, '__file__') and m.__file__ == obj_file:
if obj_module is not m:
linecache.updatecache(obj_file, m.__dict__)
def getimmediatesource(obj):
"""A variant of inspect.getsource that ignores the __wrapped__ property."""
with _linecache_lock:
_fix_linecache_record(obj)
lines, lnum = inspect.findsource(obj)
return ''.join(inspect.getblock(lines[lnum:]))
def getnamespace(f):
"""Returns the complete namespace of a function.
Namespace is defined here as the mapping of all non-local variables to values.
This includes the globals and the closure variables. Note that this captures
the entire globals collection of the function, and may contain extra symbols
that it does not actually use.
Args:
f: User defined function.
Returns:
A dict mapping symbol names to values.
"""
namespace = dict(f.__globals__)
closure = f.__closure__
freevars = f.__code__.co_freevars
if freevars and closure:
for name, cell in zip(freevars, closure):
try:
namespace[name] = cell.cell_contents
except ValueError:
# Cell contains undefined variable, omit it from the namespace.
pass
return namespace
def getqualifiedname(namespace, object_, max_depth=5, visited=None):
"""Returns the name by which a value can be referred to in a given namespace.
If the object defines a parent module, the function attempts to use it to
locate the object.
This function will recurse inside modules, but it will not search objects for
attributes. The recursion depth is controlled by max_depth.
Args:
namespace: Dict[str, Any], the namespace to search into.
object_: Any, the value to search.
max_depth: Optional[int], a limit to the recursion depth when searching
inside modules.
visited: Optional[Set[int]], ID of modules to avoid visiting.
Returns: Union[str, None], the fully-qualified name that resolves to the value
o, or None if it couldn't be found.
"""
if visited is None:
visited = set()
# Copy the dict to avoid "changed size error" during concurrent invocations.
# TODO(mdan): This is on the hot path. Can we avoid the copy?
namespace = dict(namespace)
for name in namespace:
# The value may be referenced by more than one symbol, case in which
# any symbol will be fine. If the program contains symbol aliases that
# change over time, this may capture a symbol that will later point to
# something else.
# TODO(mdan): Prefer the symbol that matches the value type name.
if object_ is namespace[name]:
return name
# If an object is not found, try to search its parent modules.
parent = tf_inspect.getmodule(object_)
if (parent is not None and parent is not object_ and parent is not namespace):
# No limit to recursion depth because of the guard above.
parent_name = getqualifiedname(
namespace, parent, max_depth=0, visited=visited)
if parent_name is not None:
name_in_parent = getqualifiedname(
parent.__dict__, object_, max_depth=0, visited=visited)
assert name_in_parent is not None, (
'An object should always be found in its owner module')
return '{}.{}'.format(parent_name, name_in_parent)
if max_depth:
# Iterating over a copy prevents "changed size due to iteration" errors.
# It's unclear why those occur - suspecting new modules may load during
# iteration.
for name in namespace.keys():
value = namespace[name]
if tf_inspect.ismodule(value) and id(value) not in visited:
visited.add(id(value))
name_in_module = getqualifiedname(value.__dict__, object_,
max_depth - 1, visited)
if name_in_module is not None:
return '{}.{}'.format(name, name_in_module)
return None
def getdefiningclass(m, owner_class):
"""Resolves the class (e.g. one of the superclasses) that defined a method."""
method_name = m.__name__
for super_class in inspect.getmro(owner_class):
if ((hasattr(super_class, '__dict__') and
method_name in super_class.__dict__) or
(hasattr(super_class, '__slots__') and
method_name in super_class.__slots__)):
return super_class
return owner_class
def getmethodclass(m):
"""Resolves a function's owner, e.g.
a method's class.
Note that this returns the object that the function was retrieved from, not
necessarily the class where it was defined.
This function relies on Python stack frame support in the interpreter, and
has the same limitations that inspect.currentframe.
Limitations. This function will only work correctly if the owned class is
visible in the caller's global or local variables.
Args:
m: A user defined function
Returns:
The class that this function was retrieved from, or None if the function
is not an object or class method, or the class that owns the object or
method is not visible to m.
Raises:
ValueError: if the class could not be resolved for any unexpected reason.
"""
# Callable objects: return their own class.
if (not hasattr(m, '__name__') and hasattr(m, '__class__') and
hasattr(m, '__call__')):
if isinstance(m.__class__, type):
return m.__class__
# Instance and class: return the class of "self".
m_self = getattr(m, '__self__', None)
if m_self is not None:
if inspect.isclass(m_self):
return m_self
return m_self.__class__
# Class, static and unbound methods: search all defined classes in any
# namespace. This is inefficient but more robust a method.
owners = []
caller_frame = tf_inspect.currentframe().f_back
try:
# TODO(mdan): This doesn't consider cell variables.
# TODO(mdan): This won't work if the owner is hidden inside a container.
# Cell variables may be pulled using co_freevars and the closure.
for v in itertools.chain(caller_frame.f_locals.values(),
caller_frame.f_globals.values()):
if hasattr(v, m.__name__):
candidate = getattr(v, m.__name__)
# Py2 methods may be bound or unbound, extract im_func to get the
# underlying function.
if hasattr(candidate, 'im_func'):
candidate = candidate.im_func
if hasattr(m, 'im_func'):
m = m.im_func
if candidate is m:
owners.append(v)
finally:
del caller_frame
if owners:
if len(owners) == 1:
return owners[0]
# If multiple owners are found, and are not subclasses, raise an error.
owner_types = tuple(o if tf_inspect.isclass(o) else type(o) for o in owners)
for o in owner_types:
if tf_inspect.isclass(o) and issubclass(o, tuple(owner_types)):
return o
raise ValueError('Found too many owners of %s: %s' % (m, owners))
return None
def getfutureimports(entity):
"""Detects what future imports are necessary to safely execute entity source.
Args:
entity: Any object
Returns:
A tuple of future strings
"""
if not (tf_inspect.isfunction(entity) or tf_inspect.ismethod(entity)):
return tuple()
return tuple(
sorted(name for name, value in entity.__globals__.items()
if getattr(value, '__module__', None) == '__future__'))
@@ -0,0 +1,612 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for inspect_utils module."""
import abc
import collections
import functools
import textwrap
import types
from tensorflow.python import lib
from tensorflow.python.autograph.pyct import inspect_utils
from tensorflow.python.autograph.pyct.testing import basic_definitions
from tensorflow.python.autograph.pyct.testing import decorators
from tensorflow.python.framework import constant_op
from tensorflow.python.platform import test
def decorator(f):
return f
def function_decorator():
def dec(f):
return f
return dec
def wrapping_decorator():
def dec(f):
def replacement(*_):
return None
@functools.wraps(f)
def wrapper(*args, **kwargs):
return replacement(*args, **kwargs)
return wrapper
return dec
class TestClass:
def member_function(self):
pass
@decorator
def decorated_member(self):
pass
@function_decorator()
def fn_decorated_member(self):
pass
@wrapping_decorator()
def wrap_decorated_member(self):
pass
@staticmethod
def static_method():
pass
@classmethod
def class_method(cls):
pass
def free_function():
pass
def factory():
return free_function
def free_factory():
def local_function():
pass
return local_function
class InspectUtilsTest(test.TestCase):
def test_islambda(self):
def test_fn():
pass
self.assertTrue(inspect_utils.islambda(lambda x: x))
self.assertFalse(inspect_utils.islambda(test_fn))
def test_islambda_renamed_lambda(self):
l = lambda x: 1
l.__name__ = 'f'
self.assertTrue(inspect_utils.islambda(l))
def test_isnamedtuple(self):
nt = collections.namedtuple('TestNamedTuple', ['a', 'b'])
class NotANamedTuple(tuple):
pass
self.assertTrue(inspect_utils.isnamedtuple(nt))
self.assertFalse(inspect_utils.isnamedtuple(NotANamedTuple))
def test_isnamedtuple_confounder(self):
"""This test highlights false positives when detecting named tuples."""
class NamedTupleLike(tuple):
_fields = ('a', 'b')
self.assertTrue(inspect_utils.isnamedtuple(NamedTupleLike))
def test_isnamedtuple_subclass(self):
"""This test highlights false positives when detecting named tuples."""
class NamedTupleSubclass(collections.namedtuple('Test', ['a', 'b'])):
pass
self.assertTrue(inspect_utils.isnamedtuple(NamedTupleSubclass))
def assertSourceIdentical(self, actual, expected):
self.assertEqual(
textwrap.dedent(actual).strip(),
textwrap.dedent(expected).strip()
)
def test_getimmediatesource_basic(self):
def test_decorator(f):
def f_wrapper(*args, **kwargs):
return f(*args, **kwargs)
return f_wrapper
expected = """
def f_wrapper(*args, **kwargs):
return f(*args, **kwargs)
"""
@test_decorator
def test_fn(a):
"""Test docstring."""
return [a]
self.assertSourceIdentical(
inspect_utils.getimmediatesource(test_fn), expected)
def test_getimmediatesource_noop_decorator(self):
def test_decorator(f):
return f
expected = '''
@test_decorator
def test_fn(a):
"""Test docstring."""
return [a]
'''
@test_decorator
def test_fn(a):
"""Test docstring."""
return [a]
self.assertSourceIdentical(
inspect_utils.getimmediatesource(test_fn), expected)
def test_getimmediatesource_functools_wrapper(self):
def wrapper_decorator(f):
@functools.wraps(f)
def wrapper(*args, **kwargs):
return f(*args, **kwargs)
return wrapper
expected = textwrap.dedent("""
@functools.wraps(f)
def wrapper(*args, **kwargs):
return f(*args, **kwargs)
""")
@wrapper_decorator
def test_fn(a):
"""Test docstring."""
return [a]
self.assertSourceIdentical(
inspect_utils.getimmediatesource(test_fn), expected)
def test_getimmediatesource_functools_wrapper_different_module(self):
expected = textwrap.dedent("""
@functools.wraps(f)
def wrapper(*args, **kwargs):
return f(*args, **kwargs)
""")
@decorators.wrapping_decorator
def test_fn(a):
"""Test docstring."""
return [a]
self.assertSourceIdentical(
inspect_utils.getimmediatesource(test_fn), expected)
def test_getimmediatesource_normal_decorator_different_module(self):
expected = textwrap.dedent("""
def standalone_wrapper(*args, **kwargs):
return f(*args, **kwargs)
""")
@decorators.standalone_decorator
def test_fn(a):
"""Test docstring."""
return [a]
self.assertSourceIdentical(
inspect_utils.getimmediatesource(test_fn), expected)
def test_getimmediatesource_normal_functional_decorator_different_module(
self):
expected = textwrap.dedent("""
def functional_wrapper(*args, **kwargs):
return f(*args, **kwargs)
""")
@decorators.functional_decorator()
def test_fn(a):
"""Test docstring."""
return [a]
self.assertSourceIdentical(
inspect_utils.getimmediatesource(test_fn), expected)
def test_getnamespace_globals(self):
ns = inspect_utils.getnamespace(factory)
self.assertEqual(ns['free_function'], free_function)
def test_getnamespace_closure_with_undefined_var(self):
if False: # pylint:disable=using-constant-test
a = 1
def test_fn():
return a
ns = inspect_utils.getnamespace(test_fn)
self.assertNotIn('a', ns)
a = 2
ns = inspect_utils.getnamespace(test_fn)
self.assertEqual(ns['a'], 2)
def test_getnamespace_hermetic(self):
# Intentionally hiding the global function to make sure we don't overwrite
# it in the global namespace.
free_function = object() # pylint:disable=redefined-outer-name
def test_fn():
return free_function
ns = inspect_utils.getnamespace(test_fn)
globs = test_fn.__globals__
self.assertTrue(ns['free_function'] is free_function)
self.assertFalse(globs['free_function'] is free_function)
def test_getnamespace_locals(self):
def called_fn():
return 0
closed_over_list = []
closed_over_primitive = 1
def local_fn():
closed_over_list.append(1)
local_var = 1
return called_fn() + local_var + closed_over_primitive
ns = inspect_utils.getnamespace(local_fn)
self.assertEqual(ns['called_fn'], called_fn)
self.assertEqual(ns['closed_over_list'], closed_over_list)
self.assertEqual(ns['closed_over_primitive'], closed_over_primitive)
self.assertTrue('local_var' not in ns)
def test_getqualifiedname(self):
foo = object()
qux = types.ModuleType('quxmodule')
bar = types.ModuleType('barmodule')
baz = object()
bar.baz = baz
ns = {
'foo': foo,
'bar': bar,
'qux': qux,
}
self.assertIsNone(inspect_utils.getqualifiedname(ns, inspect_utils))
self.assertEqual(inspect_utils.getqualifiedname(ns, foo), 'foo')
self.assertEqual(inspect_utils.getqualifiedname(ns, bar), 'bar')
self.assertEqual(inspect_utils.getqualifiedname(ns, baz), 'bar.baz')
def test_getqualifiedname_efficiency(self):
foo = object()
# We create a densely connected graph consisting of a relatively small
# number of modules and hide our symbol in one of them. The path to the
# symbol is at least 10, and each node has about 10 neighbors. However,
# by skipping visited modules, the search should take much less.
ns = {}
prev_level = []
for i in range(10):
current_level = []
for j in range(10):
mod_name = 'mod_{}_{}'.format(i, j)
mod = types.ModuleType(mod_name)
current_level.append(mod)
if i == 9 and j == 9:
mod.foo = foo
if prev_level:
# All modules at level i refer to all modules at level i+1
for prev in prev_level:
for mod in current_level:
prev.__dict__[mod.__name__] = mod
else:
for mod in current_level:
ns[mod.__name__] = mod
prev_level = current_level
self.assertIsNone(inspect_utils.getqualifiedname(ns, inspect_utils))
self.assertIsNotNone(
inspect_utils.getqualifiedname(ns, foo, max_depth=10000000000))
def test_getqualifiedname_cycles(self):
foo = object()
# We create a graph of modules that contains circular references. The
# search process should avoid them. The searched object is hidden at the
# bottom of a path of length roughly 10.
ns = {}
mods = []
for i in range(10):
mod = types.ModuleType('mod_{}'.format(i))
if i == 9:
mod.foo = foo
# Module i refers to module i+1
if mods:
mods[-1].__dict__[mod.__name__] = mod
else:
ns[mod.__name__] = mod
# Module i refers to all modules j < i.
for prev in mods:
mod.__dict__[prev.__name__] = prev
mods.append(mod)
self.assertIsNone(inspect_utils.getqualifiedname(ns, inspect_utils))
self.assertIsNotNone(
inspect_utils.getqualifiedname(ns, foo, max_depth=10000000000))
def test_getqualifiedname_finds_via_parent_module(self):
# TODO(mdan): This test is vulnerable to change in the lib module.
# A better way to forge modules should be found.
self.assertEqual(
inspect_utils.getqualifiedname(
lib.__dict__, lib.io.file_io.FileIO, max_depth=1),
'io.file_io.FileIO')
def test_getmethodclass(self):
self.assertEqual(
inspect_utils.getmethodclass(free_function), None)
self.assertEqual(
inspect_utils.getmethodclass(free_factory()), None)
self.assertEqual(
inspect_utils.getmethodclass(TestClass.member_function),
TestClass)
self.assertEqual(
inspect_utils.getmethodclass(TestClass.decorated_member),
TestClass)
self.assertEqual(
inspect_utils.getmethodclass(TestClass.fn_decorated_member),
TestClass)
self.assertEqual(
inspect_utils.getmethodclass(TestClass.wrap_decorated_member),
TestClass)
self.assertEqual(
inspect_utils.getmethodclass(TestClass.static_method),
TestClass)
self.assertEqual(
inspect_utils.getmethodclass(TestClass.class_method),
TestClass)
test_obj = TestClass()
self.assertEqual(
inspect_utils.getmethodclass(test_obj.member_function),
TestClass)
self.assertEqual(
inspect_utils.getmethodclass(test_obj.decorated_member),
TestClass)
self.assertEqual(
inspect_utils.getmethodclass(test_obj.fn_decorated_member),
TestClass)
self.assertEqual(
inspect_utils.getmethodclass(test_obj.wrap_decorated_member),
TestClass)
self.assertEqual(
inspect_utils.getmethodclass(test_obj.static_method),
TestClass)
self.assertEqual(
inspect_utils.getmethodclass(test_obj.class_method),
TestClass)
def test_getmethodclass_locals(self):
def local_function():
pass
class LocalClass:
def member_function(self):
pass
@decorator
def decorated_member(self):
pass
@function_decorator()
def fn_decorated_member(self):
pass
@wrapping_decorator()
def wrap_decorated_member(self):
pass
self.assertEqual(
inspect_utils.getmethodclass(local_function), None)
self.assertEqual(
inspect_utils.getmethodclass(LocalClass.member_function),
LocalClass)
self.assertEqual(
inspect_utils.getmethodclass(LocalClass.decorated_member),
LocalClass)
self.assertEqual(
inspect_utils.getmethodclass(LocalClass.fn_decorated_member),
LocalClass)
self.assertEqual(
inspect_utils.getmethodclass(LocalClass.wrap_decorated_member),
LocalClass)
test_obj = LocalClass()
self.assertEqual(
inspect_utils.getmethodclass(test_obj.member_function),
LocalClass)
self.assertEqual(
inspect_utils.getmethodclass(test_obj.decorated_member),
LocalClass)
self.assertEqual(
inspect_utils.getmethodclass(test_obj.fn_decorated_member),
LocalClass)
self.assertEqual(
inspect_utils.getmethodclass(test_obj.wrap_decorated_member),
LocalClass)
def test_getmethodclass_callables(self):
class TestCallable:
def __call__(self):
pass
c = TestCallable()
self.assertEqual(inspect_utils.getmethodclass(c), TestCallable)
def test_getmethodclass_no_bool_conversion(self):
tensor = constant_op.constant([1])
self.assertEqual(
inspect_utils.getmethodclass(tensor.get_shape), type(tensor))
def test_getdefiningclass(self):
class Superclass:
def foo(self):
pass
def bar(self):
pass
@classmethod
def class_method(cls):
pass
class Subclass(Superclass):
def foo(self):
pass
def baz(self):
pass
self.assertIs(
inspect_utils.getdefiningclass(Subclass.foo, Subclass), Subclass)
self.assertIs(
inspect_utils.getdefiningclass(Subclass.bar, Subclass), Superclass)
self.assertIs(
inspect_utils.getdefiningclass(Subclass.baz, Subclass), Subclass)
self.assertIs(
inspect_utils.getdefiningclass(Subclass.class_method, Subclass),
Superclass)
def test_isbuiltin(self):
self.assertTrue(inspect_utils.isbuiltin(enumerate))
self.assertTrue(inspect_utils.isbuiltin(eval))
self.assertTrue(inspect_utils.isbuiltin(float))
self.assertTrue(inspect_utils.isbuiltin(int))
self.assertTrue(inspect_utils.isbuiltin(len))
self.assertTrue(inspect_utils.isbuiltin(range))
self.assertTrue(inspect_utils.isbuiltin(zip))
self.assertFalse(inspect_utils.isbuiltin(function_decorator))
def test_isconstructor(self):
class OrdinaryClass:
pass
class OrdinaryCallableClass:
def __call__(self):
pass
class Metaclass(type):
pass
class CallableMetaclass(type):
def __call__(cls):
pass
self.assertTrue(inspect_utils.isconstructor(OrdinaryClass))
self.assertTrue(inspect_utils.isconstructor(OrdinaryCallableClass))
self.assertTrue(inspect_utils.isconstructor(Metaclass))
self.assertTrue(inspect_utils.isconstructor(Metaclass('TestClass', (), {})))
self.assertTrue(inspect_utils.isconstructor(CallableMetaclass))
self.assertFalse(inspect_utils.isconstructor(
CallableMetaclass('TestClass', (), {})))
def test_isconstructor_abc_callable(self):
class AbcBase(metaclass=abc.ABCMeta):
@abc.abstractmethod
def __call__(self):
pass
class AbcSubclass(AbcBase):
def __init__(self):
pass
def __call__(self):
pass
self.assertTrue(inspect_utils.isconstructor(AbcBase))
self.assertTrue(inspect_utils.isconstructor(AbcSubclass))
def test_getfutureimports_functions(self):
imps = inspect_utils.getfutureimports(basic_definitions.function_with_print)
self.assertNotIn('absolute_import', imps)
self.assertNotIn('division', imps)
self.assertNotIn('print_function', imps)
self.assertNotIn('generators', imps)
def test_getfutureimports_lambdas(self):
imps = inspect_utils.getfutureimports(basic_definitions.simple_lambda)
self.assertNotIn('absolute_import', imps)
self.assertNotIn('division', imps)
self.assertNotIn('print_function', imps)
self.assertNotIn('generators', imps)
def test_getfutureimports_methods(self):
imps = inspect_utils.getfutureimports(
basic_definitions.SimpleClass.method_with_print)
self.assertNotIn('absolute_import', imps)
self.assertNotIn('division', imps)
self.assertNotIn('print_function', imps)
self.assertNotIn('generators', imps)
if __name__ == '__main__':
test.main()
+19
View File
@@ -0,0 +1,19 @@
# Copyright 2017 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.
# ==============================================================================
# Test that runs inspect_utils_test as a .par file.
SCRIPT_DIR="$(dirname ${BASH_SOURCE[0]})"
${SCRIPT_DIR}/inspect_utils_test.par
+102
View File
@@ -0,0 +1,102 @@
# Copyright 2017 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.
# ==============================================================================
"""Converting AST to code and Python entities.
Adapted from Tangent.
"""
import atexit
import errno
import importlib
import os
import sys
import tempfile
from tensorflow.python.autograph.pyct import origin_info
from tensorflow.python.autograph.pyct import parser
def _remove_file(file_name):
"""Remove a file, if it exists."""
try:
os.remove(file_name)
except OSError as e:
if e.errno == errno.ENOENT:
# The file disappeared. Ignore this. Temporary files might get
# cleaned up, especially if they reside in /tmp.
pass
else:
raise
def load_source(source, delete_on_exit):
"""Loads the given source code as a Python module."""
with tempfile.NamedTemporaryFile(
mode='w',
suffix='.py',
prefix='__autograph_generated_file',
delete=False,
encoding='utf-8') as f:
module_name = os.path.basename(f.name[:-3])
file_name = f.name
f.write(source)
if delete_on_exit:
atexit.register(lambda: _remove_file(file_name))
spec = importlib.util.spec_from_file_location(module_name, file_name)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
# TODO(mdan): Use our own garbage-collected cache instead of sys.modules.
sys.modules[module_name] = module
return module, file_name
def load_ast(nodes,
indentation=' ',
include_source_map=False,
delete_on_exit=True):
"""Loads the given AST as a Python module.
Compiling the AST code this way ensures that the source code is readable by
e.g. `pdb` or `inspect`.
Args:
nodes: Union[ast.AST, Iterable[ast.AST]], the code to compile, as an AST
object.
indentation: Text, the string to use for indentation.
include_source_map: bool, whether return a source map.
delete_on_exit: bool, whether to delete the temporary file used for
compilation on exit.
Returns:
Tuple[module, Text, Dict[LineLocation, OriginInfo]], containing:
the module containing the unparsed nodes, the source code corresponding to
nodes, and the source map. Is include_source_map is False, the source map
will be None.
"""
if not isinstance(nodes, (list, tuple)):
nodes = (nodes,)
source = parser.unparse(nodes, indentation=indentation)
module, _ = load_source(source, delete_on_exit)
if include_source_map:
source_map = origin_info.create_source_map(nodes, source, module.__file__)
else:
source_map = None
# TODO(mdan): Return a structured object.
return module, source, source_map
@@ -0,0 +1,123 @@
# coding=utf-8
# Copyright 2017 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.
# ==============================================================================
"""Tests for loader module."""
import os
import textwrap
import gast
from tensorflow.python.autograph.pyct import ast_util
from tensorflow.python.autograph.pyct import loader
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import pretty_printer
from tensorflow.python.platform import test
from tensorflow.python.util import tf_inspect
class LoaderTest(test.TestCase):
def assertAstMatches(self, actual_node, expected_node_src):
expected_node = gast.parse(expected_node_src).body[0]
msg = 'AST did not match expected:\n{}\nActual:\n{}'.format(
pretty_printer.fmt(expected_node),
pretty_printer.fmt(actual_node))
self.assertTrue(ast_util.matches(actual_node, expected_node), msg)
def test_parse_load_identity(self):
def test_fn(x):
a = True
b = ''
if a:
b = (x + 1)
return b
node, _ = parser.parse_entity(test_fn, future_features=())
module, _, _ = loader.load_ast(node)
source = tf_inspect.getsource(module.test_fn)
expected_node_src = textwrap.dedent(tf_inspect.getsource(test_fn))
self.assertAstMatches(node, source)
self.assertAstMatches(node, expected_node_src)
def test_load_ast(self):
node = gast.FunctionDef(
name='f',
args=gast.arguments(
args=[
gast.Name(
'a', ctx=gast.Param(), annotation=None, type_comment=None)
],
posonlyargs=[],
vararg=None,
kwonlyargs=[],
kw_defaults=[],
kwarg=None,
defaults=[]),
body=[
gast.Return(
gast.BinOp(
op=gast.Add(),
left=gast.Name(
'a',
ctx=gast.Load(),
annotation=None,
type_comment=None),
right=gast.Constant(1, kind=None)))
],
decorator_list=[],
returns=None,
type_comment=None)
module, source, _ = loader.load_ast(node)
expected_node_src = """
# coding=utf-8
def f(a):
return (a + 1)
"""
expected_node_src = textwrap.dedent(expected_node_src)
self.assertAstMatches(node, source)
self.assertAstMatches(node, expected_node_src)
self.assertEqual(2, module.f(1))
with open(module.__file__, 'r') as temp_output:
self.assertAstMatches(node, temp_output.read())
def test_load_source(self):
test_source = textwrap.dedent(u"""
# coding=utf-8
def f(a):
'日本語 Δθₜ ← Δθₜ₋₁ + ∇Q(sₜ, aₜ)(rₜ + γₜ₊₁ max Q(⋅))'
return a + 1
""")
module, _ = loader.load_source(test_source, delete_on_exit=True)
self.assertEqual(module.f(1), 2)
self.assertEqual(
module.f.__doc__, '日本語 Δθₜ ← Δθₜ₋₁ + ∇Q(sₜ, aₜ)(rₜ + γₜ₊₁ max Q(⋅))')
def test_cleanup(self):
test_source = textwrap.dedent('')
_, filename = loader.load_source(test_source, delete_on_exit=True)
# Clean up the file before loader.py tries to remove it, to check that the
# latter can deal with that situation.
os.unlink(filename)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,53 @@
# Copyright 2017 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.
# ==============================================================================
"""Symbol naming utilities."""
from tensorflow.python.autograph.pyct import qual_names
class Namer(object):
"""Symbol name generator."""
def __init__(self, global_namespace):
self.global_namespace = global_namespace
self.generated_names = set()
def new_symbol(self, name_root, reserved_locals):
"""See control_flow.SymbolNamer.new_symbol."""
# reserved_locals may contain QNs.
all_reserved_locals = set()
for s in reserved_locals:
if isinstance(s, qual_names.QN):
all_reserved_locals.update(s.qn)
elif isinstance(s, str):
all_reserved_locals.add(s)
else:
raise ValueError('Unexpected symbol type "%s"' % type(s))
pieces = name_root.split('_')
if pieces[-1].isdigit():
name_root = '_'.join(pieces[:-1])
n = int(pieces[-1])
else:
n = 0
new_name = name_root
while (new_name in self.global_namespace or
new_name in all_reserved_locals or new_name in self.generated_names):
n += 1
new_name = '%s_%d' % (name_root, n)
self.generated_names.add(new_name)
return new_name
@@ -0,0 +1,44 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for naming module."""
from tensorflow.python.autograph.pyct import naming
from tensorflow.python.platform import test
class NamerTest(test.TestCase):
def test_new_symbol_tracks_names(self):
namer = naming.Namer({})
self.assertEqual('temp', namer.new_symbol('temp', set()))
self.assertItemsEqual(('temp',), namer.generated_names)
def test_new_symbol_avoids_duplicates(self):
namer = naming.Namer({})
self.assertEqual('temp', namer.new_symbol('temp', set()))
self.assertEqual('temp_1', namer.new_symbol('temp', set()))
self.assertItemsEqual(('temp', 'temp_1'), namer.generated_names)
def test_new_symbol_avoids_conflicts(self):
namer = naming.Namer({'temp': 1})
# temp is reserved in the global namespace
self.assertEqual('temp_1', namer.new_symbol('temp', set()))
# temp_2 is reserved in the local namespace
self.assertEqual('temp_3', namer.new_symbol('temp', set(('temp_2',))))
self.assertItemsEqual(('temp_1', 'temp_3'), namer.generated_names)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,296 @@
# Copyright 2018 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.
# ==============================================================================
"""Container for origin source code information before AutoGraph compilation."""
import collections
import difflib
import io
import os
import tokenize
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import ast_util
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import pretty_printer
from tensorflow.python.util import tf_inspect
class LineLocation(
collections.namedtuple('LineLocation', ('filename', 'lineno'))):
"""Similar to Location, but without column information.
Attributes:
filename: Text
lineno: int, 1-based
"""
pass
class Location(
collections.namedtuple('Location', ('filename', 'lineno', 'col_offset'))):
"""Encodes code location information.
Attributes:
filename: Text
lineno: int, 1-based
col_offset: int
line_loc: LineLocation
"""
@property
def line_loc(self):
return LineLocation(self.filename, self.lineno)
class OriginInfo(
collections.namedtuple(
'OriginInfo',
('loc', 'function_name', 'source_code_line', 'comment'))):
"""Container for information about the source code before conversion.
Attributes:
loc: Location
function_name: Optional[Text]
source_code_line: Text
comment: Optional[Text]
"""
def as_frame(self):
"""Returns a 4-tuple consistent with the return of traceback.extract_tb."""
return (self.loc.filename, self.loc.lineno, self.function_name,
self.source_code_line)
def __repr__(self):
if self.loc.filename:
return '{}:{}:{}'.format(
os.path.split(self.loc.filename)[1], self.loc.lineno,
self.loc.col_offset)
return '<no file>:{}:{}'.format(self.loc.lineno, self.loc.col_offset)
# TODO(mdan): This source map should be a class - easier to refer to.
def create_source_map(nodes, code, filepath):
"""Creates a source map between an annotated AST and the code it compiles to.
Note: this function assumes nodes nodes, code and filepath correspond to the
same code.
Args:
nodes: Iterable[ast.AST, ...], one or more AST modes.
code: Text, the source code in which nodes are found.
filepath: Text
Returns:
Dict[LineLocation, OriginInfo], mapping locations in code to locations
indicated by origin annotations in node.
"""
reparsed_nodes = parser.parse(code, preamble_len=0, single_node=False)
for node in reparsed_nodes:
resolve(node, code, filepath, node.lineno, node.col_offset)
source_map = {}
try:
for before, after in ast_util.parallel_walk(nodes, reparsed_nodes):
# Note: generated code might not be mapped back to its origin.
# TODO(mdan): Generated code should always be mapped to something.
origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None)
final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None)
if origin_info is None or final_info is None:
continue
# Note: the keys are by line only, excluding the column offset.
line_loc = LineLocation(final_info.loc.filename, final_info.loc.lineno)
existing_origin = source_map.get(line_loc)
if existing_origin is not None:
# Overlaps may exist because of child nodes, but almost never to
# different line locations. Exception make decorated functions, where
# both lines are mapped to the same line in the AST.
# Line overlaps: keep bottom node.
if existing_origin.loc.line_loc == origin_info.loc.line_loc:
if existing_origin.loc.lineno >= origin_info.loc.lineno:
continue
# In case of column overlaps, keep the leftmost node.
if existing_origin.loc.col_offset <= origin_info.loc.col_offset:
continue
source_map[line_loc] = origin_info
except ValueError as err:
new_msg = 'Inconsistent ASTs detected. This is a bug. Cause: \n'
new_msg += str(err)
new_msg += 'Diff:\n'
for n, rn in zip(nodes, reparsed_nodes):
nodes_str = pretty_printer.fmt(n, color=False, noanno=True)
reparsed_nodes_str = pretty_printer.fmt(rn, color=False, noanno=True)
diff = difflib.context_diff(
nodes_str.split('\n'),
reparsed_nodes_str.split('\n'),
fromfile='Original nodes',
tofile='Reparsed nodes',
n=7)
diff = '\n'.join(diff)
new_msg += diff + '\n'
raise ValueError(new_msg)
return source_map
class _Function:
def __init__(self, name):
self.name = name
class OriginResolver(gast.NodeVisitor):
"""Annotates an AST with additional source information like file name."""
def __init__(self, root_node, source_lines, comments_map,
context_lineno, context_col_offset,
filepath):
self._source_lines = source_lines
self._comments_map = comments_map
if (hasattr(root_node, 'decorator_list') and root_node.decorator_list and
hasattr(root_node.decorator_list[0], 'lineno')):
# Typical case: functions. The line number of the first decorator
# is more accurate than the line number of the function itself in
# 3.8+. In earlier versions they coincide.
self._lineno_offset = context_lineno - root_node.decorator_list[0].lineno
else:
# Fall back to the line number of the root node.
self._lineno_offset = context_lineno - root_node.lineno
self._col_offset = context_col_offset - root_node.col_offset
self._filepath = filepath
self._function_stack = []
def _absolute_lineno(self, lineno):
return lineno + self._lineno_offset
def _absolute_col_offset(self, col_offset):
if col_offset is None:
return 0
return col_offset + self._col_offset
def _attach_origin_info(self, node):
lineno = getattr(node, 'lineno', None)
col_offset = getattr(node, 'col_offset', None)
if lineno is None:
return
if self._function_stack:
function_name = self._function_stack[-1].name
else:
function_name = None
source_code_line = self._source_lines[lineno - 1]
comment = self._comments_map.get(lineno)
loc = Location(self._filepath, self._absolute_lineno(lineno),
self._absolute_col_offset(col_offset))
origin = OriginInfo(loc, function_name, source_code_line, comment)
anno.setanno(node, 'lineno', lineno)
anno.setanno(node, anno.Basic.ORIGIN, origin)
def visit(self, node):
entered_function = False
if isinstance(node, gast.FunctionDef):
entered_function = True
self._function_stack.append(_Function(node.name))
self._attach_origin_info(node)
self.generic_visit(node)
if entered_function:
self._function_stack.pop()
def resolve(node, source, context_filepath, context_lineno, context_col_offset):
"""Adds origin information to an AST, based on the source it was loaded from.
This allows us to map the original source code line numbers to generated
source code.
Note: the AST may be a part of a larger context (e.g. a function is part of
a module that may contain other things). However, this function does not
assume the source argument contains the entire context, nor that it contains
only code corresponding to node itself. However, it assumes that node was
parsed from the given source code.
For this reason, two extra arguments are required, and they indicate the
location of the node in the original context.
Args:
node: gast.AST, the AST to annotate.
source: Text, the source code representing node.
context_filepath: Text
context_lineno: int
context_col_offset: int
"""
# TODO(mdan): Pull this to a separate utility.
code_reader = io.StringIO(source)
comments_map = {}
try:
for token in tokenize.generate_tokens(code_reader.readline):
tok_type, tok_string, loc, _, _ = token
srow, _ = loc
if tok_type == tokenize.COMMENT:
comments_map[srow] = tok_string.strip()[1:].strip()
except tokenize.TokenError:
if isinstance(node, gast.Lambda):
# Source code resolution in older Python versions is brittle for
# lambda functions, and may contain garbage.
pass
else:
raise
source_lines = source.split('\n')
visitor = OriginResolver(node, source_lines, comments_map,
context_lineno, context_col_offset,
context_filepath)
visitor.visit(node)
def resolve_entity(node, source, entity):
"""Like resolve, but extracts the context information from an entity."""
lines, lineno = tf_inspect.getsourcelines(entity)
filepath = tf_inspect.getsourcefile(entity)
# Poor man's attempt at guessing the column offset: count the leading
# whitespace. This might not work well with tabs.
definition_line = lines[0]
col_offset = len(definition_line) - len(definition_line.lstrip())
resolve(node, source, filepath, lineno, col_offset)
def copy_origin(from_node, to_node):
"""Copies the origin info from a node to another, recursively."""
origin = anno.Basic.ORIGIN.of(from_node, default=None)
if origin is None:
return
if not isinstance(to_node, (list, tuple)):
to_node = (to_node,)
for node in to_node:
for n in gast.walk(node):
anno.setanno(n, anno.Basic.ORIGIN, origin)
@@ -0,0 +1,268 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for origin_info module."""
import inspect
import sys
import textwrap
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import inspect_utils
from tensorflow.python.autograph.pyct import origin_info
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct.testing import basic_definitions
from tensorflow.python.platform import test
from tensorflow.python.util import tf_inspect
class OriginInfoTest(test.TestCase):
def test_create_source_map(self):
source = """
def test_fn(x):
return x + 1
"""
source = textwrap.dedent(source)
node = parser.parse(source)
fake_origin = origin_info.OriginInfo(
loc=origin_info.Location('fake_filename', 3, 7),
function_name='fake_function_name',
source_code_line='fake source line',
comment=None)
anno.setanno(node, anno.Basic.ORIGIN, fake_origin)
source_map = origin_info.create_source_map(node, source, 'test_filename')
loc = origin_info.LineLocation('test_filename', 2)
self.assertIn(loc, source_map)
self.assertIs(source_map[loc], fake_origin)
def _create_source_map(self, test_fn):
node, source = parser.parse_entity(test_fn, ())
origin_info.resolve_entity(node, source, test_fn)
# Creating a source map with the source code as output will create
# an identity map.
return origin_info.create_source_map(node, source, 'test_filename')
def test_create_source_map_identity(self):
test_fn = basic_definitions.simple_function
source_map = self._create_source_map(test_fn)
module_path = tf_inspect.getsourcefile(test_fn)
# Origin line numbers below should match those in basic_definitions.py
fn_start = inspect.getsourcelines(test_fn)[1]
definition_loc = origin_info.LineLocation('test_filename', 1)
self.assertIn(definition_loc, source_map)
self.assertEqual(source_map[definition_loc].loc.lineno, fn_start)
self.assertEqual(source_map[definition_loc].loc.filename, module_path)
self.assertEqual(source_map[definition_loc].function_name,
'simple_function')
def test_create_source_map_multiline_call(self):
test_fn = basic_definitions.function_with_multiline_call
source_map = self._create_source_map(test_fn)
module_path = tf_inspect.getsourcefile(test_fn)
# Origin line numbers below should match those in basic_definitions.py
fn_start = inspect.getsourcelines(test_fn)[1]
call_loc = origin_info.LineLocation('test_filename', 3)
self.assertIn(call_loc, source_map)
self.assertEqual(source_map[call_loc].loc.lineno, fn_start + 2)
self.assertEqual(source_map[call_loc].loc.filename, module_path)
self.assertEqual(source_map[call_loc].function_name,
'function_with_multiline_call')
self.assertEqual(source_map[call_loc].source_code_line, ' return range(')
second_arg_loc = origin_info.LineLocation('test_filename', 5)
self.assertIn(second_arg_loc, source_map)
self.assertEqual(source_map[second_arg_loc].loc.lineno, fn_start + 4)
self.assertEqual(source_map[second_arg_loc].loc.filename, module_path)
self.assertEqual(source_map[second_arg_loc].function_name,
'function_with_multiline_call')
self.assertEqual(source_map[second_arg_loc].source_code_line,
' x + 1,')
def test_create_source_map_no_origin_info(self):
test_fn = basic_definitions.simple_function
node, _ = parser.parse_entity(test_fn,
inspect_utils.getfutureimports(test_fn))
# No origin information should result in an empty map.
test_fn_lines, _ = tf_inspect.getsourcelines(test_fn)
source_map = origin_info.create_source_map(node, '\n'.join(test_fn_lines),
test_fn)
self.assertEmpty(source_map)
def test_resolve(self):
source = """
def test_fn(x):
'''Docstring.'''
return x # comment
"""
source = textwrap.dedent(source)
node = parser.parse(source)
origin_info.resolve(node, source, 'test_file', 10, 10)
def_origin = anno.getanno(node, anno.Basic.ORIGIN)
self.assertEqual(def_origin.loc.filename, 'test_file')
self.assertEqual(def_origin.loc.lineno, 10)
self.assertEqual(def_origin.loc.col_offset, 10)
self.assertEqual(def_origin.source_code_line, 'def test_fn(x):')
self.assertIsNone(def_origin.comment)
docstring_origin = anno.getanno(node.body[0], anno.Basic.ORIGIN)
self.assertEqual(def_origin.loc.filename, 'test_file')
self.assertEqual(docstring_origin.loc.lineno, 11)
self.assertEqual(docstring_origin.loc.col_offset, 12)
self.assertEqual(docstring_origin.source_code_line, " '''Docstring.'''")
self.assertIsNone(docstring_origin.comment)
ret_origin = anno.getanno(node.body[1], anno.Basic.ORIGIN)
self.assertEqual(def_origin.loc.filename, 'test_file')
self.assertEqual(ret_origin.loc.lineno, 12)
self.assertEqual(ret_origin.loc.col_offset, 12)
self.assertEqual(ret_origin.source_code_line, ' return x # comment')
self.assertEqual(ret_origin.comment, 'comment')
def test_resolve_with_trailing_garbage(self):
# This comment will be missed because the tokenizer fails to reach it.
source = ' lambda: foo([], bar=1)), baz=2)()'
clean_source = 'lambda: foo([], bar=1)'
node = parser.parse(clean_source).value
origin_info.resolve(node, source, 'test_file', 10, 10)
def_origin = anno.getanno(node, anno.Basic.ORIGIN)
self.assertEqual(def_origin.loc.lineno, 10)
self.assertEqual(def_origin.loc.col_offset, 10)
self.assertEqual(def_origin.source_code_line, source)
self.assertIsNone(def_origin.comment)
def test_resolve_entity(self):
test_fn = basic_definitions.simple_function
node, source = parser.parse_entity(
test_fn, inspect_utils.getfutureimports(test_fn))
origin_info.resolve_entity(node, source, test_fn)
# The line numbers below should match those in basic_definitions.py
fn_start = inspect.getsourcelines(test_fn)[1]
def_origin = anno.getanno(node, anno.Basic.ORIGIN)
self.assertEqual(def_origin.loc.lineno, fn_start)
self.assertEqual(def_origin.loc.col_offset, 0)
self.assertEqual(def_origin.source_code_line, 'def simple_function(x):')
self.assertIsNone(def_origin.comment)
docstring_origin = anno.getanno(node.body[0], anno.Basic.ORIGIN)
self.assertEqual(docstring_origin.loc.lineno, fn_start + 1)
self.assertEqual(docstring_origin.loc.col_offset, 2)
self.assertEqual(docstring_origin.source_code_line, ' """Docstring."""')
self.assertIsNone(docstring_origin.comment)
ret_origin = anno.getanno(node.body[1], anno.Basic.ORIGIN)
self.assertEqual(ret_origin.loc.lineno, fn_start + 2)
self.assertEqual(ret_origin.loc.col_offset, 2)
self.assertEqual(ret_origin.source_code_line, ' return x # comment')
self.assertEqual(ret_origin.comment, 'comment')
def test_resolve_entity_nested_function(self):
test_fn = basic_definitions.nested_functions
node, source = parser.parse_entity(
test_fn, inspect_utils.getfutureimports(test_fn))
origin_info.resolve_entity(node, source, test_fn)
# The line numbers below should match those in basic_definitions.py
fn_start = inspect.getsourcelines(test_fn)[1]
inner_def_origin = anno.getanno(node.body[1], anno.Basic.ORIGIN)
self.assertEqual(inner_def_origin.loc.lineno, fn_start + 3)
self.assertEqual(inner_def_origin.loc.col_offset, 2)
self.assertEqual(inner_def_origin.source_code_line, ' def inner_fn(y):')
self.assertIsNone(inner_def_origin.comment)
inner_ret_origin = anno.getanno(node.body[1].body[0], anno.Basic.ORIGIN)
self.assertEqual(inner_ret_origin.loc.lineno, fn_start + 4)
self.assertEqual(inner_ret_origin.loc.col_offset, 4)
self.assertEqual(inner_ret_origin.source_code_line, ' return y')
self.assertIsNone(inner_ret_origin.comment)
def test_resolve_entity_indented_block(self):
test_fn = basic_definitions.SimpleClass.simple_method
node, source = parser.parse_entity(test_fn,
inspect_utils.getfutureimports(test_fn))
origin_info.resolve_entity(node, source, test_fn)
# The line numbers below should match those in basic_definitions.py
fn_start = inspect.getsourcelines(test_fn)[1]
def_origin = anno.getanno(node, anno.Basic.ORIGIN)
self.assertEqual(def_origin.loc.lineno, fn_start)
self.assertEqual(def_origin.loc.col_offset, 2)
self.assertEqual(def_origin.source_code_line, 'def simple_method(self):')
self.assertIsNone(def_origin.comment)
ret_origin = anno.getanno(node.body[0], anno.Basic.ORIGIN)
self.assertEqual(ret_origin.loc.lineno, fn_start + 1)
self.assertEqual(ret_origin.loc.col_offset, 4)
self.assertEqual(ret_origin.source_code_line, ' return self')
self.assertIsNone(ret_origin.comment)
def test_resolve_entity_decorated_function(self):
test_fn = basic_definitions.decorated_function
node, source = parser.parse_entity(test_fn,
inspect_utils.getfutureimports(test_fn))
origin_info.resolve_entity(node, source, test_fn)
# The line numbers below should match those in basic_definitions.py
fn_start = inspect.getsourcelines(test_fn)[1]
def_origin = anno.getanno(node, anno.Basic.ORIGIN)
if sys.version_info >= (3, 8):
self.assertEqual(def_origin.loc.lineno, fn_start + 2)
self.assertEqual(def_origin.source_code_line,
'def decorated_function(x):')
else:
self.assertEqual(def_origin.loc.lineno, fn_start)
self.assertEqual(def_origin.source_code_line, '@basic_decorator')
self.assertEqual(def_origin.loc.col_offset, 0)
self.assertIsNone(def_origin.comment)
if_origin = anno.getanno(node.body[0], anno.Basic.ORIGIN)
self.assertEqual(if_origin.loc.lineno, fn_start + 3)
self.assertEqual(if_origin.loc.col_offset, 2)
self.assertEqual(if_origin.source_code_line, ' if x > 0:')
self.assertIsNone(if_origin.comment)
ret1_origin = anno.getanno(node.body[0].body[0], anno.Basic.ORIGIN)
self.assertEqual(ret1_origin.loc.lineno, fn_start + 4)
self.assertEqual(ret1_origin.loc.col_offset, 4)
self.assertEqual(ret1_origin.source_code_line, ' return 1')
self.assertIsNone(ret1_origin.comment)
ret2_origin = anno.getanno(node.body[1], anno.Basic.ORIGIN)
self.assertEqual(ret2_origin.loc.lineno, fn_start + 5)
self.assertEqual(ret2_origin.loc.col_offset, 2)
self.assertEqual(ret2_origin.source_code_line, ' return 2')
self.assertIsNone(ret2_origin.comment)
if __name__ == '__main__':
test.main()
+391
View File
@@ -0,0 +1,391 @@
# Copyright 2017 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.
# ==============================================================================
"""Converting code to AST.
Adapted from Tangent.
"""
import ast
import inspect
import io
import linecache
import re
import sys
import textwrap
import tokenize
import gast
from tensorflow.python.autograph.pyct import errors
from tensorflow.python.autograph.pyct import inspect_utils
from tensorflow.python.util import tf_inspect
PY2_PREAMBLE = textwrap.dedent("""
""")
PY3_PREAMBLE = ''
MAX_SIZE = 0
if sys.version_info >= (3,):
STANDARD_PREAMBLE = PY3_PREAMBLE
MAX_SIZE = sys.maxsize
else:
STANDARD_PREAMBLE = PY2_PREAMBLE
MAX_SIZE = sys.maxint
STANDARD_PREAMBLE_LEN = STANDARD_PREAMBLE.count('__future__')
_LEADING_WHITESPACE = re.compile(r'\s*')
def _unfold_continuations(code_string):
"""Removes any backslash line continuations from the code."""
return code_string.replace('\\\n', '')
def dedent_block(code_string):
"""Dedents a code so that its first line starts at row zero."""
code_string = _unfold_continuations(code_string)
token_gen = tokenize.generate_tokens(io.StringIO(code_string).readline)
block_indentation = None
tokens = []
try:
for tok in token_gen:
tokens.append(tok)
except tokenize.TokenError:
# Resolution of lambda functions may yield incomplete code, which can
# in turn generate this error. We silently ignore this error because the
# parser may still be able to deal with it.
pass
for tok in tokens:
tok_type, tok_string, _, _, _ = tok
if tok_type == tokenize.INDENT:
block_indentation = tok_string
block_level = len(block_indentation)
break
elif tok_type not in (
tokenize.NL, tokenize.NEWLINE, tokenize.STRING, tokenize.COMMENT):
block_indentation = ''
break
if not block_indentation:
return code_string
block_level = len(block_indentation)
first_indent_uses_tabs = '\t' in block_indentation
for i, tok in enumerate(tokens):
tok_type, tok_string, _, _, _ = tok
if tok_type == tokenize.INDENT:
if ((' ' in tok_string and first_indent_uses_tabs)
or ('\t' in tok_string and not first_indent_uses_tabs)):
# TODO(mdan): We could attempt to convert tabs to spaces by unix rule.
# See:
# https://docs.python.org/3/reference/lexical_analysis.html#indentation
raise errors.UnsupportedLanguageElementError(
'code mixing tabs and spaces for indentation is not allowed')
if len(tok_string) >= block_level:
tok_string = tok_string[block_level:]
tokens[i] = (tok_type, tok_string)
new_code = tokenize.untokenize(tokens)
# Note: untokenize respects the line structure, but not the whitespace within
# lines. For example, `def foo()` may be untokenized as `def foo ()`
# So instead of using the output of dedent, we match the leading whitespace
# on each line.
dedented_code = []
for line, new_line in zip(code_string.split('\n'), new_code.split('\n')):
original_indent = re.match(_LEADING_WHITESPACE, line).group()
new_indent = re.match(_LEADING_WHITESPACE, new_line).group()
if len(original_indent) > len(new_indent):
dedented_line = line[len(original_indent) - len(new_indent):]
else:
dedented_line = line
dedented_code.append(dedented_line)
new_code = '\n'.join(dedented_code)
return new_code
def parse_entity(entity, future_features):
"""Returns the AST and source code of given entity.
Args:
entity: Any, Python function/method/class
future_features: Iterable[Text], future features to use (e.g.
'print_statement'). See
https://docs.python.org/2/reference/simple_stmts.html#future
Returns:
gast.AST, Text: the parsed AST node; the source code that was parsed to
generate the AST (including any prefixes that this function may have added).
"""
if inspect_utils.islambda(entity):
return _parse_lambda(entity)
try:
original_source = inspect_utils.getimmediatesource(entity)
except OSError as e:
raise errors.InaccessibleSourceCodeError(
f'Unable to locate the source code of {entity}. Note that functions'
' defined in certain environments, like the interactive Python shell,'
' do not expose their source code. If that is the case, you should'
' define them in a .py source file. If you are certain the code is'
' graph-compatible, wrap the call using'
f' @tf.autograph.experimental.do_not_convert. Original error: {e}')
source = dedent_block(original_source)
future_statements = tuple(
'from __future__ import {}'.format(name) for name in future_features)
source = '\n'.join(future_statements + (source,))
return parse(source, preamble_len=len(future_features)), source
def _without_context(node, lines, minl, maxl):
"""Returns a clean node and source code without indenting and context."""
for n in gast.walk(node):
lineno = getattr(n, 'lineno', None)
if lineno is not None:
n.lineno = lineno - minl
end_lineno = getattr(n, 'end_lineno', None)
if end_lineno is not None:
n.end_lineno = end_lineno - minl
code_lines = lines[minl - 1:maxl]
# Attempt to clean up surrounding context code.
end_col_offset = getattr(node, 'end_col_offset', None)
if end_col_offset is not None:
# This is only available in 3.8.
code_lines[-1] = code_lines[-1][:end_col_offset]
col_offset = getattr(node, 'col_offset', None)
if col_offset is None:
# Older Python: try to find the "lambda" token. This is brittle.
match = re.search(r'(?<!\w)lambda(?!\w)', code_lines[0])
if match is not None:
col_offset = match.start(0)
if col_offset is not None:
code_lines[0] = code_lines[0][col_offset:]
code_block = '\n'.join([c.rstrip() for c in code_lines])
return node, code_block
def _arg_name(node):
if node is None:
return None
if isinstance(node, gast.Name):
return node.id
assert isinstance(node, str)
return node
def _node_matches_argspec(node, func):
"""Returns True is node fits the argspec of func."""
# TODO(mdan): Use just inspect once support for Python 2 is dropped.
arg_spec = tf_inspect.getfullargspec(func)
node_args = tuple(_arg_name(arg) for arg in node.args.args)
if node_args != tuple(arg_spec.args):
return False
if arg_spec.varargs != _arg_name(node.args.vararg):
return False
if arg_spec.varkw != _arg_name(node.args.kwarg):
return False
node_kwonlyargs = tuple(_arg_name(arg) for arg in node.args.kwonlyargs)
if node_kwonlyargs != tuple(arg_spec.kwonlyargs):
return False
return True
def _parse_lambda(lam):
"""Returns the AST and source code of given lambda function.
Args:
lam: types.LambdaType, Python function/method/class
Returns:
gast.AST, Text: the parsed AST node; the source code that was parsed to
generate the AST (including any prefixes that this function may have added).
"""
# TODO(mdan): Use a fast path if the definition is not multi-line.
# We could detect that the lambda is in a multi-line expression by looking
# at the surrounding code - an surrounding set of parentheses indicates a
# potential multi-line definition.
mod = inspect.getmodule(lam)
f = inspect.getsourcefile(lam)
def_line = lam.__code__.co_firstlineno
# This method is more robust that just calling inspect.getsource(mod), as it
# works in interactive shells, where getsource would fail. This is the
# same procedure followed by inspect for non-modules:
# https://github.com/python/cpython/blob/3.8/Lib/inspect.py#L772
lines = linecache.getlines(f, mod.__dict__)
source = ''.join(lines)
# Narrow down to the last node starting before our definition node.
all_nodes = parse(source, preamble_len=0, single_node=False)
search_nodes = []
for node in all_nodes:
# Also include nodes without a line number, for safety. This is defensive -
# we don't know whether such nodes might exist, and if they do, whether
# they are not safe to skip.
# TODO(mdan): Replace this check with an assertion or skip such nodes.
if getattr(node, 'lineno', def_line) <= def_line:
search_nodes.append(node)
else:
# Found a node starting past our lambda - can stop the search.
break
# Extract all lambda nodes from the shortlist.
lambda_nodes = []
for node in search_nodes:
lambda_nodes.extend(
n for n in gast.walk(node) if isinstance(n, gast.Lambda))
# Filter down to lambda nodes which span our actual lambda.
candidates = []
for ln in lambda_nodes:
minl, maxl = MAX_SIZE, 0
for n in gast.walk(ln):
minl = min(minl, getattr(n, 'lineno', minl))
lineno = getattr(n, 'lineno', maxl)
end_lineno = getattr(n, 'end_lineno', None)
if end_lineno is not None:
# end_lineno is more precise, but lineno should almost always work too.
lineno = end_lineno
maxl = max(maxl, lineno)
if minl <= def_line <= maxl:
candidates.append((ln, minl, maxl))
# Happy path: exactly one node found.
if len(candidates) == 1:
(node, minl, maxl), = candidates # pylint:disable=unbalanced-tuple-unpacking
return _without_context(node, lines, minl, maxl)
elif not candidates:
lambda_codes = '\n'.join([unparse(l) for l in lambda_nodes])
raise errors.UnsupportedLanguageElementError(
f'could not parse the source code of {lam}:'
f' no matching AST found among candidates:\n{lambda_codes}')
# Attempt to narrow down selection by signature is multiple nodes are found.
matches = [v for v in candidates if _node_matches_argspec(v[0], lam)]
if len(matches) == 1:
(node, minl, maxl), = matches
return _without_context(node, lines, minl, maxl)
# Give up if could not narrow down to a single node.
matches = '\n'.join(
'Match {}:\n{}\n'.format(i, unparse(node, include_encoding_marker=False))
for i, (node, _, _) in enumerate(matches))
raise errors.UnsupportedLanguageElementError(
f'could not parse the source code of {lam}: found multiple definitions'
' with identical signatures at the location. This error'
' may be avoided by defining each lambda on a single line and with'
f' unique argument names. The matching definitions were:\n{matches}')
# TODO(mdan): This should take futures as input instead.
def parse(src, preamble_len=0, single_node=True):
"""Returns the AST of given piece of code.
Args:
src: Text
preamble_len: Int, indicates leading nodes in the parsed AST which should be
dropped.
single_node: Bool, whether `src` is assumed to be represented by exactly one
AST node.
Returns:
ast.AST
"""
module_node = gast.parse(src)
nodes = module_node.body
if preamble_len:
nodes = nodes[preamble_len:]
if single_node:
if len(nodes) != 1:
raise ValueError('expected exactly one node, got {}'.format(nodes))
return nodes[0]
return nodes
def parse_expression(src):
"""Returns the AST of given identifier.
Args:
src: A piece of code that represents a single Python expression
Returns:
A gast.AST object.
Raises:
ValueError: if src does not consist of a single Expression.
"""
src = STANDARD_PREAMBLE + src.strip()
node = parse(src, preamble_len=STANDARD_PREAMBLE_LEN, single_node=True)
if __debug__:
if not isinstance(node, gast.Expr):
raise ValueError(
'expected exactly one node of type Expr, got {}'.format(node))
return node.value
def unparse(node, indentation=None, include_encoding_marker=True):
"""Returns the source code of given AST.
Args:
node: The code to compile, as an AST object.
indentation: Unused, deprecated. The returning code will always be indented
at 4 spaces.
include_encoding_marker: Bool, whether to include a comment on the first
line to explicitly specify UTF-8 encoding.
Returns:
code: The source code generated from the AST object
source_mapping: A mapping between the user and AutoGraph generated code.
"""
del indentation # astunparse doesn't allow configuring it.
if not isinstance(node, (list, tuple)):
node = (node,)
codes = []
if include_encoding_marker:
codes.append('# coding=utf-8')
for n in node:
if isinstance(n, gast.AST):
ast_n = gast.gast_to_ast(n)
else:
ast_n = n
ast.fix_missing_locations(ast_n)
codes.append(ast.unparse(ast_n).strip())
return '\n'.join(codes)
@@ -0,0 +1,390 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for parser module."""
import re
import sys
import textwrap
import gast
from tensorflow.python.autograph.pyct import ast_util
from tensorflow.python.autograph.pyct import errors
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import pretty_printer
from tensorflow.python.platform import test
class ParserTest(test.TestCase):
def assertAstMatches(self, actual_node, expected_node_src, expr=True):
if expr:
# Ensure multi-line expressions parse.
expected_node = gast.parse('({})'.format(expected_node_src)).body[0]
expected_node = expected_node.value
else:
expected_node = gast.parse(expected_node_src).body[0]
msg = 'AST did not match expected:\n{}\nActual:\n{}'.format(
pretty_printer.fmt(expected_node),
pretty_printer.fmt(actual_node))
self.assertTrue(ast_util.matches(actual_node, expected_node), msg)
def test_parse_entity(self):
def f(x):
return x + 1
node, _ = parser.parse_entity(f, future_features=())
self.assertEqual('f', node.name)
def test_parse_lambda(self):
l = lambda x: x + 1
expected_node_src = 'lambda x: (x + 1)'
node, source = parser.parse_entity(l, future_features=())
self.assertAstMatches(node, source)
self.assertAstMatches(node, expected_node_src)
def test_parse_lambda_prefix_cleanup(self):
lambda_lam = lambda x: x + 1
expected_node_src = 'lambda x: (x + 1)'
node, source = parser.parse_entity(lambda_lam, future_features=())
self.assertAstMatches(node, source)
self.assertAstMatches(node, expected_node_src)
def test_parse_lambda_resolution_by_location(self):
_ = lambda x: x + 1
l = lambda x: x + 1
_ = lambda x: x + 1
expected_node_src = 'lambda x: (x + 1)'
node, source = parser.parse_entity(l, future_features=())
self.assertAstMatches(node, source)
self.assertAstMatches(node, expected_node_src)
self.assertEqual(source, 'lambda x: x + 1')
def test_parse_lambda_resolution_by_signature(self):
l = lambda x: lambda x, y: x + y
node, source = parser.parse_entity(l, future_features=())
expected_node_src = 'lambda x: (lambda x, y: (x + y))'
self.assertAstMatches(node, source)
self.assertAstMatches(node, expected_node_src)
self.assertEqual(source, 'lambda x: lambda x, y: x + y')
node, source = parser.parse_entity(l(0), future_features=())
expected_node_src = 'lambda x, y: (x + y)'
self.assertAstMatches(node, source)
self.assertAstMatches(node, expected_node_src)
self.assertEqual(source, 'lambda x, y: x + y')
def test_parse_lambda_resolution_ambiguous(self):
l = lambda x: lambda x: 2 * x
expected_exception_text = re.compile(r'found multiple definitions'
r'.+'
r'\(?lambda x: \(?lambda x'
r'.+'
r'\(?lambda x: \(?2', re.DOTALL)
with self.assertRaisesRegex(
errors.UnsupportedLanguageElementError,
expected_exception_text):
parser.parse_entity(l, future_features=())
with self.assertRaisesRegex(
errors.UnsupportedLanguageElementError,
expected_exception_text):
parser.parse_entity(l(0), future_features=())
def assertMatchesWithPotentialGarbage(self, source, expected, garbage):
# In runtimes which don't track end_col_number, the source contains the
# entire line, which in turn may have garbage from the surrounding context.
self.assertIn(source, (expected, expected + garbage))
def test_parse_lambda_multiline(self):
l = (
lambda x: lambda y: x + y # pylint:disable=g-long-lambda
- 1)
node, source = parser.parse_entity(l, future_features=())
expected_node_src = 'lambda x: (lambda y: ((x + y) - 1))'
self.assertAstMatches(node, expected_node_src)
self.assertMatchesWithPotentialGarbage(
source, ('lambda x: lambda y: x + y # pylint:disable=g-long-lambda\n'
' - 1'), ')')
node, source = parser.parse_entity(l(0), future_features=())
expected_node_src = 'lambda y: ((x + y) - 1)'
self.assertAstMatches(node, expected_node_src)
self.assertMatchesWithPotentialGarbage(
source, ('lambda y: x + y # pylint:disable=g-long-lambda\n'
' - 1'), ')')
def test_parse_lambda_in_expression(self):
l = (
lambda x: lambda y: x + y + 1,
lambda x: lambda y: x + y + 2,
)
node, source = parser.parse_entity(l[0], future_features=())
expected_node_src = 'lambda x: (lambda y: ((x + y) + 1))'
self.assertAstMatches(node, expected_node_src)
self.assertMatchesWithPotentialGarbage(
source, 'lambda x: lambda y: x + y + 1', ',')
node, source = parser.parse_entity(l[0](0), future_features=())
expected_node_src = 'lambda y: ((x + y) + 1)'
self.assertAstMatches(node, expected_node_src)
self.assertMatchesWithPotentialGarbage(
source, 'lambda y: x + y + 1', ',')
node, source = parser.parse_entity(l[1], future_features=())
expected_node_src = 'lambda x: (lambda y: ((x + y) + 2))'
self.assertAstMatches(node, expected_node_src)
self.assertMatchesWithPotentialGarbage(source,
'lambda x: lambda y: x + y + 2', ',')
node, source = parser.parse_entity(l[1](0), future_features=())
expected_node_src = 'lambda y: ((x + y) + 2)'
self.assertAstMatches(node, expected_node_src)
self.assertMatchesWithPotentialGarbage(source, 'lambda y: x + y + 2', ',')
def test_parse_lambda_complex_body(self):
l = lambda x: ( # pylint:disable=g-long-lambda
x.y(
[],
x.z,
(),
x[0:2],
),
x.u,
'abc',
1,
)
node, source = parser.parse_entity(l, future_features=())
expected_node_src = "lambda x: (x.y([], x.z, (), x[0:2]), x.u, 'abc', 1)"
self.assertAstMatches(node, expected_node_src)
base_source = ('lambda x: ( # pylint:disable=g-long-lambda\n'
' x.y(\n'
' [],\n'
' x.z,\n'
' (),\n'
' x[0:2],\n'
' ),\n'
' x.u,\n'
' \'abc\',\n'
' 1,')
# The complete source includes the trailing parenthesis. But that is only
# detected in runtimes which correctly track end_lineno for ASTs.
self.assertMatchesWithPotentialGarbage(source, base_source, '\n )')
def test_parse_lambda_function_call_definition(self):
def do_parse_and_test(lam, **unused_kwargs):
node, source = parser.parse_entity(lam, future_features=())
expected_node_src = 'lambda x: x'
self.assertAstMatches(node, expected_node_src)
self.assertMatchesWithPotentialGarbage(
source, 'lambda x: x', ', named_arg=1)')
do_parse_and_test( # Intentional line break
lambda x: x, named_arg=1)
def test_parse_entity_print_function(self):
def f(x):
print(x)
node, _ = parser.parse_entity(f, future_features=('print_function',))
self.assertEqual('f', node.name)
def test_parse_comments(self):
def f():
# unindented comment
pass
node, _ = parser.parse_entity(f, future_features=())
self.assertEqual('f', node.name)
def test_parse_multiline_strings(self):
def f():
print("""
multiline
string""")
node, _ = parser.parse_entity(f, future_features=())
self.assertEqual('f', node.name)
def _eval_code(self, code, name):
globs = {}
exec(code, globs) # pylint:disable=exec-used
return globs[name]
def test_dedent_block_basic(self):
code = """
def f(x):
if x > 0:
return -x
return x
"""
f = self._eval_code(parser.dedent_block(code), 'f')
self.assertEqual(f(1), -1)
self.assertEqual(f(-1), -1)
def test_dedent_block_comments_out_of_line(self):
code = """
###
def f(x):
###
if x > 0:
###
return -x
###
###
return x
###
"""
f = self._eval_code(parser.dedent_block(code), 'f')
self.assertEqual(f(1), -1)
self.assertEqual(f(-1), -1)
def test_dedent_block_multiline_string(self):
code = """
def f():
'''
Docstring.
'''
return '''
1
2
3'''
"""
f = self._eval_code(parser.dedent_block(code), 'f')
if sys.version_info >= (3, 13):
self.assertEqual(f.__doc__.strip(), 'Docstring.')
else:
self.assertEqual(f.__doc__, '\n Docstring.\n ')
self.assertEqual(f(), '\n 1\n 2\n 3')
def test_dedent_block_multiline_expression(self):
code = """
def f():
return (1,
2,
3)
"""
f = self._eval_code(parser.dedent_block(code), 'f')
self.assertEqual(f(), (1, 2, 3))
def test_dedent_block_continuation(self):
code = r"""
def f():
a = \
1
return a
"""
f = self._eval_code(parser.dedent_block(code), 'f')
self.assertEqual(f(), 1)
def test_dedent_block_continuation_in_string(self):
code = r"""
def f():
a = "a \
b"
return a
"""
f = self._eval_code(parser.dedent_block(code), 'f')
self.assertEqual(f(), 'a b')
def test_parse_expression(self):
node = parser.parse_expression('a.b')
self.assertEqual('a', node.value.id)
self.assertEqual('b', node.attr)
def test_unparse(self):
node = gast.If(
test=gast.Constant(1, kind=None),
body=[
gast.Assign(
targets=[
gast.Name(
'a',
ctx=gast.Store(),
annotation=None,
type_comment=None)
],
value=gast.Name(
'b', ctx=gast.Load(), annotation=None, type_comment=None))
],
orelse=[
gast.Assign(
targets=[
gast.Name(
'a',
ctx=gast.Store(),
annotation=None,
type_comment=None)
],
value=gast.Constant('c', kind=None))
])
source = parser.unparse(node, indentation=' ')
self.assertEqual(
textwrap.dedent("""
# coding=utf-8
if 1:
a = b
else:
a = 'c'
""").strip(), source.strip())
def test_ext_slice_roundtrip(self):
def ext_slice(n):
return n[:, :], n[0, :], n[:, 0]
node, _ = parser.parse_entity(ext_slice, future_features=())
source = parser.unparse(node)
self.assertAstMatches(node, source, expr=False)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,130 @@
# Copyright 2017 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.
# ==============================================================================
"""Print an AST tree in a form more readable than ast.dump."""
import gast
import termcolor
class PrettyPrinter(gast.NodeVisitor):
"""Print AST nodes."""
def __init__(self, color, noanno):
self.indent_lvl = 0
self.result = ''
self.color = color
self.noanno = noanno
def _color(self, string, color, attrs=None):
if self.color:
return termcolor.colored(string, color, attrs=attrs)
return string
def _type(self, node):
return self._color(node.__class__.__name__, None, ['bold'])
def _field(self, name):
return self._color(name, 'blue')
def _value(self, name):
return self._color(name, 'magenta')
def _warning(self, name):
return self._color(name, 'red')
def _indent(self):
return self._color('| ' * self.indent_lvl, None, ['dark'])
def _print(self, s):
self.result += s
self.result += '\n'
def generic_visit(self, node, name=None):
# In very rare instances, a list can contain something other than a Node.
# e.g. Global contains a list of strings.
if isinstance(node, str):
if name:
self._print('%s%s="%s"' % (self._indent(), name, node))
else:
self._print('%s"%s"' % (self._indent(), node))
return
if node._fields:
cont = ':'
else:
cont = '()'
if name:
self._print('%s%s=%s%s' % (self._indent(), self._field(name),
self._type(node), cont))
else:
self._print('%s%s%s' % (self._indent(), self._type(node), cont))
self.indent_lvl += 1
for f in node._fields:
if self.noanno and f.startswith('__'):
continue
if not hasattr(node, f):
self._print('%s%s' % (self._indent(), self._warning('%s=<unset>' % f)))
continue
v = getattr(node, f)
if isinstance(v, list):
if v:
self._print('%s%s=[' % (self._indent(), self._field(f)))
self.indent_lvl += 1
for n in v:
if n is not None:
self.generic_visit(n)
else:
self._print('%sNone' % (self._indent()))
self.indent_lvl -= 1
self._print('%s]' % (self._indent()))
else:
self._print('%s%s=[]' % (self._indent(), self._field(f)))
elif isinstance(v, tuple):
if v:
self._print('%s%s=(' % (self._indent(), self._field(f)))
self.indent_lvl += 1
for n in v:
if n is not None:
self.generic_visit(n)
else:
self._print('%sNone' % (self._indent()))
self.indent_lvl -= 1
self._print('%s)' % (self._indent()))
else:
self._print('%s%s=()' % (self._indent(), self._field(f)))
elif isinstance(v, gast.AST):
self.generic_visit(v, f)
elif isinstance(v, bytes):
self._print('%s%s=%s' % (self._indent(), self._field(f),
self._value('b"%s"' % v)))
elif isinstance(v, str):
self._print('%s%s=%s' % (self._indent(), self._field(f),
self._value('u"%s"' % v)))
else:
self._print('%s%s=%s' % (self._indent(), self._field(f),
self._value(v)))
self.indent_lvl -= 1
def fmt(node, color=True, noanno=False):
printer = PrettyPrinter(color, noanno)
if isinstance(node, (list, tuple)):
for n in node:
printer.visit(n)
else:
printer.visit(node)
return printer.result
@@ -0,0 +1,61 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for pretty_printer module."""
import ast
import textwrap
from tensorflow.python.autograph.pyct import pretty_printer
from tensorflow.python.platform import test
class PrettyPrinterTest(test.TestCase):
def test_unicode_bytes(self):
source = textwrap.dedent('''
def f():
return b'b', u'u', 'depends_py2_py3'
''')
node = ast.parse(source)
self.assertIsNotNone(pretty_printer.fmt(node))
def test_format(self):
node = ast.FunctionDef(
name='f',
args=ast.arguments(
args=[ast.Name(id='a', ctx=ast.Param())],
vararg=None,
kwarg=None,
defaults=[],
),
body=[
ast.Return(
ast.BinOp(
op=ast.Add(),
left=ast.Name(id='a', ctx=ast.Load()),
right=ast.Constant(1),
)
)
],
decorator_list=[],
returns=None,
)
# Just checking for functionality, the color control characters make it
# difficult to inspect the result.
self.assertIsNotNone(pretty_printer.fmt(node))
if __name__ == '__main__':
test.main()
@@ -0,0 +1,266 @@
# Copyright 2017 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.
# ==============================================================================
"""Utilities for manipulating qualified names.
A qualified name is a uniform way to refer to simple (e.g. 'foo') and composite
(e.g. 'foo.bar') syntactic symbols.
This is *not* related to the __qualname__ attribute used by inspect, which
refers to scopes.
"""
import collections
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import parser
class CallerMustSetThis(object):
pass
class Symbol(collections.namedtuple('Symbol', ['name'])):
"""Represents a Python symbol."""
class Literal(collections.namedtuple('Literal', ['value'])):
"""Represents a Python numeric literal."""
def __str__(self):
if isinstance(self.value, str):
return "'{}'".format(self.value)
return str(self.value)
def __repr__(self):
return str(self)
# TODO(mdan): Use subclasses to remove the has_attr has_subscript booleans.
class QN(object):
"""Represents a qualified name."""
def __init__(self, base, attr=None, subscript=None):
if attr is not None and subscript is not None:
raise ValueError('A QN can only be either an attr or a subscript, not '
'both: attr={}, subscript={}.'.format(attr, subscript))
self._has_attr = False
self._has_subscript = False
if attr is not None:
if not isinstance(base, QN):
raise ValueError(
'for attribute QNs, base must be a QN; got instead "%s"' % base)
if not isinstance(attr, str):
raise ValueError('attr may only be a string; got instead "%s"' % attr)
self._parent = base
# TODO(mdan): Get rid of the tuple - it can only have 1 or 2 elements now.
self.qn = (base, attr)
self._has_attr = True
elif subscript is not None:
if not isinstance(base, QN):
raise ValueError('For subscript QNs, base must be a QN.')
self._parent = base
self.qn = (base, subscript)
self._has_subscript = True
else:
if not isinstance(base, (str, Literal)):
# TODO(mdan): Require Symbol instead of string.
raise ValueError(
'for simple QNs, base must be a string or a Literal object;'
' got instead "%s"' % type(base))
assert '.' not in base and '[' not in base and ']' not in base
self._parent = None
self.qn = (base,)
def is_symbol(self):
return isinstance(self.qn[0], str)
def is_simple(self):
return len(self.qn) <= 1
def is_composite(self):
return len(self.qn) > 1
def has_subscript(self):
return self._has_subscript
def has_attr(self):
return self._has_attr
@property
def attr(self):
if not self._has_attr:
raise ValueError('Cannot get attr of non-attribute "%s".' % self)
return self.qn[1]
@property
def parent(self):
if self._parent is None:
raise ValueError('Cannot get parent of simple name "%s".' % self.qn[0])
return self._parent
@property
def owner_set(self):
"""Returns all the symbols (simple or composite) that own this QN.
In other words, if this symbol was modified, the symbols in the owner set
may also be affected.
Examples:
'a.b[c.d]' has two owners, 'a' and 'a.b'
"""
owners = set()
if self.has_attr() or self.has_subscript():
owners.add(self.parent)
owners.update(self.parent.owner_set)
return owners
@property
def support_set(self):
"""Returns the set of simple symbols that this QN relies on.
This would be the smallest set of symbols necessary for the QN to
statically resolve (assuming properties and index ranges are verified
at runtime).
Examples:
'a.b' has only one support symbol, 'a'
'a[i]' has two support symbols, 'a' and 'i'
"""
# TODO(mdan): This might be the set of Name nodes in the AST. Track those?
roots = set()
if self.has_attr():
roots.update(self.parent.support_set)
elif self.has_subscript():
roots.update(self.parent.support_set)
roots.update(self.qn[1].support_set)
else:
roots.add(self)
return roots
def __hash__(self):
return hash(self.qn + (self._has_attr, self._has_subscript))
def __eq__(self, other):
return (isinstance(other, QN) and self.qn == other.qn and
self.has_subscript() == other.has_subscript() and
self.has_attr() == other.has_attr())
def __lt__(self, other):
return str(self) < str(other)
def __gt__(self, other):
return str(self) > str(other)
def __str__(self):
root = self.qn[0]
if self.has_subscript():
return '{}[{}]'.format(root, self.qn[1])
if self.has_attr():
return '.'.join(map(str, self.qn))
else:
return str(root)
def __repr__(self):
return str(self)
def ssf(self):
"""Simple symbol form."""
ssfs = [n.ssf() if isinstance(n, QN) else n for n in self.qn]
ssf_string = ''
for i in range(0, len(self.qn) - 1):
if self.has_subscript():
delimiter = '_sub_'
else:
delimiter = '_'
ssf_string += ssfs[i] + delimiter
return ssf_string + ssfs[-1]
def ast(self):
"""AST representation."""
# The caller must adjust the context appropriately.
if self.has_subscript():
return gast.Subscript(
value=self.parent.ast(),
slice=self.qn[-1].ast(),
ctx=CallerMustSetThis)
if self.has_attr():
return gast.Attribute(
value=self.parent.ast(), attr=self.qn[-1], ctx=CallerMustSetThis)
base = self.qn[0]
if isinstance(base, str):
return gast.Name(
base, ctx=CallerMustSetThis, annotation=None, type_comment=None)
elif isinstance(base, Literal):
return gast.Constant(base.value, kind=None)
else:
assert False, ('the constructor should prevent types other than '
'str and Literal')
class QnResolver(gast.NodeTransformer):
"""Annotates nodes with QN information.
Note: Not using NodeAnnos to avoid circular dependencies.
"""
def visit_Name(self, node):
node = self.generic_visit(node)
anno.setanno(node, anno.Basic.QN, QN(node.id))
return node
def visit_Attribute(self, node):
node = self.generic_visit(node)
if anno.hasanno(node.value, anno.Basic.QN):
anno.setanno(node, anno.Basic.QN,
QN(anno.getanno(node.value, anno.Basic.QN), attr=node.attr))
return node
def visit_Subscript(self, node):
# TODO(mdan): This may no longer apply if we overload getitem.
node = self.generic_visit(node)
s = node.slice
if isinstance(s, (gast.Tuple, gast.Slice)):
# TODO(mdan): Support range and multi-dimensional indices.
# Continuing silently because some demos use these.
return node
if isinstance(s, gast.Constant) and s.value != Ellipsis:
subscript = QN(Literal(s.value))
else:
# The index may be an expression, case in which a name doesn't make sense.
if anno.hasanno(s, anno.Basic.QN):
subscript = anno.getanno(s, anno.Basic.QN)
else:
return node
if anno.hasanno(node.value, anno.Basic.QN):
anno.setanno(node, anno.Basic.QN,
QN(anno.getanno(node.value, anno.Basic.QN),
subscript=subscript))
return node
def resolve(node):
return QnResolver().visit(node)
def from_str(qn_str):
node = parser.parse_expression(qn_str)
node = resolve(node)
return anno.getanno(node, anno.Basic.QN)
@@ -0,0 +1,261 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for qual_names module."""
import textwrap
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.autograph.pyct.qual_names import QN
from tensorflow.python.autograph.pyct.qual_names import resolve
from tensorflow.python.platform import test
class QNTest(test.TestCase):
def test_from_str(self):
a = QN('a')
b = QN('b')
a_dot_b = QN(a, attr='b')
a_sub_b = QN(a, subscript=b)
self.assertEqual(qual_names.from_str('a.b'), a_dot_b)
self.assertEqual(qual_names.from_str('a'), a)
self.assertEqual(qual_names.from_str('a[b]'), a_sub_b)
def test_basic(self):
a = QN('a')
self.assertEqual(a.qn, ('a',))
self.assertEqual(str(a), 'a')
self.assertEqual(a.ssf(), 'a')
self.assertEqual(a.ast().id, 'a')
self.assertFalse(a.is_composite())
with self.assertRaises(ValueError):
_ = a.parent
a_b = QN(a, attr='b')
self.assertEqual(a_b.qn, (a, 'b'))
self.assertEqual(str(a_b), 'a.b')
self.assertEqual(a_b.ssf(), 'a_b')
self.assertEqual(a_b.ast().value.id, 'a')
self.assertEqual(a_b.ast().attr, 'b')
self.assertTrue(a_b.is_composite())
self.assertEqual(a_b.parent.qn, ('a',))
def test_subscripts(self):
a = QN('a')
b = QN('b')
a_sub_b = QN(a, subscript=b)
self.assertEqual(a_sub_b.qn, (a, b))
self.assertEqual(str(a_sub_b), 'a[b]')
self.assertEqual(a_sub_b.ssf(), 'a_sub_b')
self.assertEqual(a_sub_b.ast().value.id, 'a')
self.assertEqual(a_sub_b.ast().slice.id, 'b')
self.assertTrue(a_sub_b.is_composite())
self.assertTrue(a_sub_b.has_subscript())
self.assertEqual(a_sub_b.parent.qn, ('a',))
c = QN('c')
b_sub_c = QN(b, subscript=c)
a_sub_b_sub_c = QN(a, subscript=b_sub_c)
self.assertEqual(a_sub_b_sub_c.qn, (a, b_sub_c))
self.assertTrue(a_sub_b_sub_c.is_composite())
self.assertTrue(a_sub_b_sub_c.has_subscript())
self.assertEqual(b_sub_c.qn, (b, c))
self.assertEqual(str(a_sub_b_sub_c), 'a[b[c]]')
self.assertEqual(a_sub_b_sub_c.ssf(), 'a_sub_b_sub_c')
self.assertEqual(a_sub_b_sub_c.ast().value.id, 'a')
self.assertEqual(a_sub_b_sub_c.ast().slice.value.id, 'b')
self.assertEqual(a_sub_b_sub_c.ast().slice.slice.id, 'c')
self.assertEqual(b_sub_c.ast().slice.id, 'c')
self.assertEqual(a_sub_b_sub_c.parent.qn, ('a',))
with self.assertRaises(ValueError):
QN('a', 'b')
def test_equality(self):
a = QN('a')
a2 = QN('a')
a_b = QN(a, attr='b')
self.assertEqual(a2.qn, ('a',))
with self.assertRaises(ValueError):
_ = a.parent
a_b2 = QN(a, attr='b')
self.assertEqual(a_b2.qn, (a, 'b'))
self.assertEqual(a_b2.parent.qn, ('a',))
self.assertTrue(a2 == a)
self.assertFalse(a2 is a)
self.assertTrue(a_b.parent == a)
self.assertTrue(a_b2.parent == a)
self.assertTrue(a_b2 == a_b)
self.assertFalse(a_b2 is a_b)
self.assertFalse(a_b2 == a)
a_sub_b = QN(a, subscript='b')
a_sub_b2 = QN(a, subscript='b')
self.assertTrue(a_sub_b == a_sub_b2)
self.assertFalse(a_sub_b == a_b)
def test_nested_attrs_subscripts(self):
a = QN('a')
b = QN('b')
c = QN('c')
b_sub_c = QN(b, subscript=c)
a_sub_b_sub_c = QN(a, subscript=b_sub_c)
b_dot_c = QN(b, attr='c')
a_sub__b_dot_c = QN(a, subscript=b_dot_c)
a_sub_b = QN(a, subscript=b)
a_sub_b__dot_c = QN(a_sub_b, attr='c')
a_dot_b = QN(a, attr='b')
a_dot_b_sub_c = QN(a_dot_b, subscript=c)
self.assertEqual(str(a_sub_b_sub_c), 'a[b[c]]')
self.assertEqual(str(a_sub__b_dot_c), 'a[b.c]')
self.assertEqual(str(a_sub_b__dot_c), 'a[b].c')
self.assertEqual(str(a_dot_b_sub_c), 'a.b[c]')
self.assertNotEqual(a_sub_b_sub_c, a_sub__b_dot_c)
self.assertNotEqual(a_sub_b_sub_c, a_sub_b__dot_c)
self.assertNotEqual(a_sub_b_sub_c, a_dot_b_sub_c)
self.assertNotEqual(a_sub__b_dot_c, a_sub_b__dot_c)
self.assertNotEqual(a_sub__b_dot_c, a_dot_b_sub_c)
self.assertNotEqual(a_sub_b__dot_c, a_dot_b_sub_c)
def test_hashable(self):
d = {QN('a'): 'a', QN('b'): 'b'}
self.assertEqual(d[QN('a')], 'a')
self.assertEqual(d[QN('b')], 'b')
self.assertNotIn(QN('c'), d)
def test_literals(self):
a = QN('a')
a_sub_str_b = QN(a, subscript=QN(qual_names.Literal('b')))
a_sub_b = QN(a, subscript=QN('b'))
self.assertNotEqual(a_sub_str_b, a_sub_b)
self.assertNotEqual(hash(a_sub_str_b), hash(a_sub_b))
self.assertEqual(a_sub_str_b.ast().slice.value, 'b')
self.assertEqual(str(a_sub_str_b), "a['b']")
a_sub_three = QN(a, subscript=QN(qual_names.Literal(3)))
self.assertEqual(a_sub_three.ast().slice.value, 3)
self.assertEqual(str(a_sub_three), 'a[3]')
def test_support_set(self):
a = QN('a')
b = QN('b')
c = QN('c')
a_sub_b = QN(a, subscript=b)
a_dot_b = QN(a, attr='b')
a_dot_b_dot_c = QN(a_dot_b, attr='c')
a_dot_b_sub_c = QN(a_dot_b, subscript=c)
self.assertSetEqual(a.support_set, set((a,)))
self.assertSetEqual(a_sub_b.support_set, set((a, b)))
self.assertSetEqual(a_dot_b.support_set, set((a,)))
self.assertSetEqual(a_dot_b_dot_c.support_set, set((a,)))
self.assertSetEqual(a_dot_b_sub_c.support_set, set((a, c)))
def test_comparison(self):
less_than_apos = chr(ord('\'') - 1)
self.assertGreater(QN('z'), QN(qual_names.Literal('a')))
self.assertLess(QN(less_than_apos), QN(qual_names.Literal('a')))
self.assertGreater(QN(qual_names.Literal('z')), QN(less_than_apos))
self.assertLess(QN(qual_names.Literal('a')), QN('z'))
class QNResolverTest(test.TestCase):
def assertQNStringIs(self, node, qn_str):
self.assertEqual(str(anno.getanno(node, anno.Basic.QN)), qn_str)
def test_resolve(self):
samples = """
a
a.b
(c, d.e)
[f, (g.h.i)]
j(k, l)
"""
nodes = parser.parse(textwrap.dedent(samples), single_node=False)
nodes = tuple(resolve(node).value for node in nodes)
self.assertQNStringIs(nodes[0], 'a')
self.assertQNStringIs(nodes[1], 'a.b')
self.assertQNStringIs(nodes[2].elts[0], 'c')
self.assertQNStringIs(nodes[2].elts[1], 'd.e')
self.assertQNStringIs(nodes[3].elts[0], 'f')
self.assertQNStringIs(nodes[3].elts[1], 'g.h.i')
self.assertQNStringIs(nodes[4].func, 'j')
self.assertQNStringIs(nodes[4].args[0], 'k')
self.assertQNStringIs(nodes[4].args[1], 'l')
def test_subscript_resolve(self):
samples = """
x[i]
x[i.b]
a.b[c]
a.b[x.y]
a[z[c]]
a[b[c[d]]]
a[b].c
a.b.c[d].e.f
a.b[c[d]].e.f
a.b[c[d.e.f].g].h
"""
nodes = parser.parse(textwrap.dedent(samples), single_node=False)
nodes = tuple(resolve(node).value for node in nodes)
self.assertQNStringIs(nodes[0], 'x[i]')
self.assertQNStringIs(nodes[1], 'x[i.b]')
self.assertQNStringIs(nodes[2], 'a.b[c]')
self.assertQNStringIs(nodes[3], 'a.b[x.y]')
self.assertQNStringIs(nodes[4], 'a[z[c]]')
self.assertQNStringIs(nodes[5], 'a[b[c[d]]]')
self.assertQNStringIs(nodes[6], 'a[b].c')
self.assertQNStringIs(nodes[7], 'a.b.c[d].e.f')
self.assertQNStringIs(nodes[8], 'a.b[c[d]].e.f')
self.assertQNStringIs(nodes[9], 'a.b[c[d.e.f].g].h')
def test_function_calls(self):
samples = """
a.b
a.b()
a().b
z[i]
z[i]()
z()[i]
"""
nodes = parser.parse(textwrap.dedent(samples), single_node=False)
nodes = tuple(resolve(node).value for node in nodes)
self.assertQNStringIs(nodes[0], 'a.b')
self.assertQNStringIs(nodes[1].func, 'a.b')
self.assertQNStringIs(nodes[2].value.func, 'a')
self.assertQNStringIs(nodes[3], 'z[i]')
self.assertQNStringIs(nodes[4].func, 'z[i]')
self.assertQNStringIs(nodes[5].value.func, 'z')
if __name__ == '__main__':
test.main()
@@ -0,0 +1,177 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
py_library(
name = "liveness",
srcs = ["liveness.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":annos",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:transformer",
"@pypi//gast",
],
)
py_library(
name = "reaching_fndefs",
srcs = ["reaching_fndefs.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:transformer",
"@pypi//gast",
],
)
py_test(
name = "reaching_fndefs_test",
srcs = ["reaching_fndefs_test.py"],
strict_deps = True,
deps = [
":activity",
":reaching_definitions",
":reaching_fndefs",
#internal proto upb dep
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:naming",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:transformer",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "activity",
srcs = ["activity.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":annos",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:transformer",
"@pypi//gast",
],
)
py_library(
name = "type_inference",
srcs = ["type_inference.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":activity",
":annos",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:transformer",
"@pypi//gast",
],
)
py_library(
name = "reaching_definitions",
srcs = ["reaching_definitions.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:transformer",
"@pypi//gast",
],
)
py_library(
name = "annos",
srcs = ["annos.py"],
strict_deps = True,
visibility = ["//visibility:public"],
)
py_test(
name = "activity_test",
srcs = ["activity_test.py"],
strict_deps = True,
deps = [
":activity",
":annos",
"@pypi//gast",
#internal proto upb dep
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:naming",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:transformer",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "liveness_test",
testonly = True,
srcs = ["liveness_test.py"],
strict_deps = True,
deps = [
":activity",
":liveness",
":reaching_fndefs",
#internal proto upb dep
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:naming",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:transformer",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "reaching_definitions_test",
srcs = ["reaching_definitions_test.py"],
strict_deps = True,
deps = [
":activity",
":reaching_definitions",
#internal proto upb dep
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:naming",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:transformer",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "type_inference_test",
srcs = ["type_inference_test.py"],
strict_deps = True,
deps = [
":activity",
":reaching_definitions",
":reaching_fndefs",
":type_inference",
#internal proto upb dep
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:transpiler",
"//tensorflow/python/platform:client_testlib",
],
)
@@ -0,0 +1,28 @@
# Copyright 2017 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.
# ==============================================================================
"""Static information resolution.
This module contains utilities to help annotate AST nodes with as much runtime
information as can be possibly extracted without actually executing the code,
under that assumption that the context in which the code will run is known.
Overall, the different analyses have the functions listed below:
* activity: inventories symbols read, written to, params, etc. at different
levels
* liveness, reaching_definitions: dataflow analyses based on the program's CFG
and using the symbol information gathered by activity analysis
"""
@@ -0,0 +1,706 @@
# Copyright 2017 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.
# ==============================================================================
"""Activity analysis.
Requires qualified name annotations (see qual_names.py).
"""
import copy
import weakref
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.autograph.pyct import transformer
from tensorflow.python.autograph.pyct.static_analysis.annos import NodeAnno
class Scope(object):
"""Encloses local symbol definition and usage information.
This can track for instance whether a symbol is modified in the current scope.
Note that scopes do not necessarily align with Python's scopes. For example,
the body of an if statement may be considered a separate scope.
Caution - the AST references held by this object are weak.
Scope objects are mutable during construction only, and must be frozen using
`Scope.finalize()` before use. Furthermore, a scope is consistent only after
all its children have been frozen. While analysing code blocks, scopes are
being gradually built, from the innermost scope outward. Freezing indicates
that the analysis of a code block is complete. Once frozen, mutation is no
longer allowed. `is_final` tracks whether the scope is frozen or not. Certain
properties, like `referenced`, are only accurate when called on frozen scopes.
Attributes:
parent: Optional[Scope], the parent scope, if any.
isolated: bool, whether the scope is a true Python scope (e.g. the scope of
a function), or just a surrogate tracking an ordinary code block. Using
the terminology of the Python 3 reference documentation, True roughly
represents an actual scope, whereas False represents an ordinary code
block.
function_name: Optional[str], name of the function owning this scope.
isolated_names: Set[qual_names.QN], identifiers that are isolated to this
scope (even if the scope is not isolated).
annotations: Set[qual_names.QN], identifiers used as type annotations
in this scope.
read: Set[qual_names.QN], identifiers read in this scope.
modified: Set[qual_names.QN], identifiers modified in this scope.
deleted: Set[qual_names.QN], identifiers deleted in this scope.
bound: Set[qual_names.QN], names that are bound to this scope. See
https://docs.python.org/3/reference/executionmodel.html#binding-of-names
for a precise definition.
globals: Set[qual_names.QN], names that are explicitly marked as global in
this scope. Note that this doesn't include free read-only vars bound to
global symbols.
nonlocals: Set[qual_names.QN], names that are explicitly marked as nonlocal
in this scope. Note that this doesn't include free read-only vars bound to
global symbols.
free_vars: Set[qual_names.QN], the free variables in this scope. See
https://docs.python.org/3/reference/executionmodel.html for a precise
definition.
params: WeakValueDictionary[qual_names.QN, ast.Node], function arguments
visible in this scope, mapped to the function node that defines them.
enclosing_scope: Scope, the innermost isolated scope that is a transitive
parent of this scope. May be the scope itself.
referenced: Set[qual_names.QN], the totality of the symbols used by this
scope and its parents.
is_final: bool, whether the scope is frozen or not.
Note - simple statements may never delete and modify a symbol at the same
time. However, compound ones like if statements can. In that latter case, it's
undefined whether the symbol is actually modified or deleted upon statement
exit. Certain analyses like reaching definitions need to be careful about
this.
"""
# Note: this mutable-immutable pattern is used because using a builder would
# have taken a lot more boilerplate.
def __init__(self, parent, isolated=True, function_name=None):
"""Create a new scope.
Args:
parent: A Scope or None.
isolated: Whether the scope is isolated, that is, whether variables
modified in this scope should be considered modified in the parent
scope.
function_name: Name of the function owning this scope.
"""
self.parent = parent
self.isolated = isolated
self.function_name = function_name
self.isolated_names = set()
self.read = set()
self.modified = set()
self.deleted = set()
self.bound = set()
self.globals = set()
self.nonlocals = set()
self.annotations = set()
self.params = weakref.WeakValueDictionary()
# Certain fields can only be accessed after the scope and all its parent
# scopes have been fully built. This field guards that.
self.is_final = False
@property
def enclosing_scope(self):
assert self.is_final
if self.parent is not None and not self.isolated:
return self.parent
return self
@property
def referenced(self):
if self.parent is not None:
return self.read | self.parent.referenced
return self.read
@property
def free_vars(self):
enclosing_scope = self.enclosing_scope
return enclosing_scope.read - enclosing_scope.bound
def copy_from(self, other):
"""Recursively copies the contents of this scope from another scope."""
assert not self.is_final
if self.parent is not None:
assert other.parent is not None
self.parent.copy_from(other.parent)
self.isolated_names = copy.copy(other.isolated_names)
self.modified = copy.copy(other.modified)
self.read = copy.copy(other.read)
self.deleted = copy.copy(other.deleted)
self.bound = copy.copy(other.bound)
self.annotations = copy.copy(other.annotations)
self.params = copy.copy(other.params)
@classmethod
def copy_of(cls, other):
if other.parent is not None:
assert other.parent is not None
parent = cls.copy_of(other.parent)
else:
parent = None
new_copy = cls(parent)
new_copy.copy_from(other)
return new_copy
def merge_from(self, other):
"""Adds all activity from another scope to this scope."""
assert not self.is_final
if self.parent is not None:
assert other.parent is not None
self.parent.merge_from(other.parent)
self.isolated_names.update(other.isolated_names)
self.read.update(other.read)
self.modified.update(other.modified)
self.bound.update(other.bound)
self.deleted.update(other.deleted)
self.annotations.update(other.annotations)
self.params.update(other.params)
def finalize(self):
"""Freezes this scope."""
assert not self.is_final
# TODO(mdan): freeze read, modified, bound.
if self.parent is not None:
assert not self.parent.is_final
if not self.isolated:
self.parent.read.update(self.read - self.isolated_names)
self.parent.modified.update(self.modified - self.isolated_names)
self.parent.bound.update(self.bound - self.isolated_names)
self.parent.globals.update(self.globals)
self.parent.nonlocals.update(self.nonlocals)
self.parent.annotations.update(self.annotations)
else:
# TODO(mdan): This is not accurate.
self.parent.read.update(self.read - self.bound)
self.parent.annotations.update(self.annotations - self.bound)
self.is_final = True
def __repr__(self):
return 'Scope{r=%s, w=%s}' % (tuple(self.read), tuple(self.modified))
def mark_param(self, name, owner):
# Assumption: all AST nodes have the same life span. This lets us use
# a weak reference to mark the connection between a symbol node and the
# function node whose argument that symbol is.
self.params[name] = owner
class _Comprehension(object):
no_root = True
def __init__(self):
# TODO(mdan): Consider using an enum.
self.is_list_comp = False
self.targets = set()
class _FunctionOrClass(object):
def __init__(self):
self.node = None
class ActivityAnalyzer(transformer.Base):
"""Annotates nodes with local scope information.
See Scope.
The use of this class requires that qual_names.resolve() has been called on
the node. This class will ignore nodes have not been
annotated with their qualified names.
"""
def __init__(self, context, parent_scope=None):
super(ActivityAnalyzer, self).__init__(context)
self.allow_skips = False
self.scope = Scope(parent_scope, isolated=True)
# Note: all these flags crucially rely on the respective nodes are
# leaves in the AST, that is, they cannot contain other statements.
self._in_aug_assign = False
self._in_annotation = False
self._track_annotations_only = False
@property
def _in_constructor(self):
context = self.state[_FunctionOrClass]
if context.level > 2:
innermost = context.stack[-1].node
parent = context.stack[-2].node
return (isinstance(parent, gast.ClassDef) and
(isinstance(innermost, gast.FunctionDef) and
innermost.name == '__init__'))
return False
def _node_sets_self_attribute(self, node):
if anno.hasanno(node, anno.Basic.QN):
qn = anno.getanno(node, anno.Basic.QN)
# TODO(mdan): The 'self' argument is not guaranteed to be called 'self'.
if qn.has_attr and qn.parent.qn == ('self',):
return True
return False
def _track_symbol(self, node, composite_writes_alter_parent=False):
if self._track_annotations_only and not self._in_annotation:
return
# A QN may be missing when we have an attribute (or subscript) on a function
# call. Example: a().b
if not anno.hasanno(node, anno.Basic.QN):
return
qn = anno.getanno(node, anno.Basic.QN)
# When inside a comprehension, ignore reads to any of the comprehensions's
# targets. This includes attributes or slices of those arguments.
for l in self.state[_Comprehension]:
if qn in l.targets:
return
if qn.owner_set & set(l.targets):
return
if isinstance(node.ctx, gast.Store):
# In comprehensions, modified symbols are the comprehension targets.
if self.state[_Comprehension].level > 0:
self.state[_Comprehension].targets.add(qn)
return
self.scope.modified.add(qn)
self.scope.bound.add(qn)
if qn.is_composite and composite_writes_alter_parent:
self.scope.modified.add(qn.parent)
if self._in_aug_assign:
self.scope.read.add(qn)
elif isinstance(node.ctx, gast.Load):
self.scope.read.add(qn)
if self._in_annotation:
self.scope.annotations.add(qn)
elif isinstance(node.ctx, gast.Param):
self.scope.bound.add(qn)
self.scope.mark_param(qn, self.state[_FunctionOrClass].node)
elif isinstance(node.ctx, gast.Del):
# The read matches the Python semantics - attempting to delete an
# undefined symbol is illegal.
self.scope.read.add(qn)
# Targets of del are considered bound:
# https://docs.python.org/3/reference/executionmodel.html#binding-of-names
self.scope.bound.add(qn)
self.scope.deleted.add(qn)
else:
raise ValueError('Unknown context {} for node "{}".'.format(
type(node.ctx), qn))
def _enter_scope(self, isolated, f_name=None):
self.scope = Scope(self.scope, isolated=isolated, function_name=f_name)
def _exit_scope(self):
exited_scope = self.scope
exited_scope.finalize()
self.scope = exited_scope.parent
return exited_scope
def _exit_and_record_scope(self, node, tag=anno.Static.SCOPE):
node_scope = self._exit_scope()
anno.setanno(node, tag, node_scope)
return node_scope
def _process_statement(self, node):
self._enter_scope(False)
node = self.generic_visit(node)
self._exit_and_record_scope(node)
return node
def _process_annotation(self, node):
self._in_annotation = True
node = self.visit(node)
self._in_annotation = False
return node
def visit_Import(self, node):
return self._process_statement(node)
def visit_ImportFrom(self, node):
return self._process_statement(node)
def visit_Global(self, node):
self._enter_scope(False)
for name in node.names:
qn = qual_names.QN(name)
self.scope.read.add(qn)
self.scope.globals.add(qn)
self._exit_and_record_scope(node)
return node
def visit_Nonlocal(self, node):
self._enter_scope(False)
for name in node.names:
qn = qual_names.QN(name)
self.scope.read.add(qn)
self.scope.bound.add(qn)
self.scope.nonlocals.add(qn)
self._exit_and_record_scope(node)
return node
def visit_Expr(self, node):
return self._process_statement(node)
def visit_Raise(self, node):
return self._process_statement(node)
def visit_Return(self, node):
return self._process_statement(node)
def visit_Assign(self, node):
return self._process_statement(node)
def visit_AnnAssign(self, node):
self._enter_scope(False)
node.target = self.visit(node.target)
if node.value is not None:
# Can be None for pure declarations, e.g. `n: int`. This is a new thing
# enabled by type annotations, but does not influence static analysis
# (declarations are not definitions).
node.value = self.visit(node.value)
if node.annotation:
node.annotation = self._process_annotation(node.annotation)
self._exit_and_record_scope(node)
return node
def visit_AugAssign(self, node):
# Special rules for AugAssign. Here, the AST only shows the target as
# written, when it is in fact also read.
self._enter_scope(False)
self._in_aug_assign = True
node.target = self.visit(node.target)
self._in_aug_assign = False
node.op = self.visit(node.op)
node.value = self.visit(node.value)
self._exit_and_record_scope(node)
return node
def visit_Delete(self, node):
return self._process_statement(node)
def visit_Name(self, node):
if node.annotation:
node.annotation = self._process_annotation(node.annotation)
self._track_symbol(node)
return node
def visit_alias(self, node):
node = self.generic_visit(node)
if node.asname is None:
# Only the root name is a real symbol operation.
qn = qual_names.QN(node.name.split('.')[0])
else:
qn = qual_names.QN(node.asname)
self.scope.modified.add(qn)
self.scope.bound.add(qn)
return node
def visit_Attribute(self, node):
node = self.generic_visit(node)
if self._in_constructor and self._node_sets_self_attribute(node):
self._track_symbol(node, composite_writes_alter_parent=True)
else:
self._track_symbol(node)
return node
def visit_Subscript(self, node):
node = self.generic_visit(node)
# Subscript writes (e.g. a[b] = "value") are considered to modify
# both the element itself (a[b]) and its parent (a).
self._track_symbol(node)
return node
def visit_Print(self, node):
self._enter_scope(False)
node.values = self.visit_block(node.values)
node_scope = self._exit_and_record_scope(node)
anno.setanno(node, NodeAnno.ARGS_SCOPE, node_scope)
return node
def visit_Assert(self, node):
return self._process_statement(node)
def visit_Call(self, node):
self._enter_scope(False)
node.args = self.visit_block(node.args)
node.keywords = self.visit_block(node.keywords)
# TODO(mdan): Account starargs, kwargs
self._exit_and_record_scope(node, tag=NodeAnno.ARGS_SCOPE)
node.func = self.visit(node.func)
return node
def _process_block_node(self, node, block, scope_name):
self._enter_scope(False)
block = self.visit_block(block)
self._exit_and_record_scope(node, tag=scope_name)
return node
def _process_parallel_blocks(self, parent, children):
# Because the scopes are not isolated, processing any child block
# modifies the parent state causing the other child blocks to be
# processed incorrectly. So we need to checkpoint the parent scope so that
# each child sees the same context.
before_parent = Scope.copy_of(self.scope)
after_children = []
for child, scope_name in children:
self.scope.copy_from(before_parent)
parent = self._process_block_node(parent, child, scope_name)
after_child = Scope.copy_of(self.scope)
after_children.append(after_child)
for after_child in after_children:
self.scope.merge_from(after_child)
return parent
def _process_comprehension(self,
node,
is_list_comp=False,
is_dict_comp=False):
with self.state[_Comprehension] as comprehension_:
comprehension_.is_list_comp = is_list_comp
# Note: it's important to visit the generators first to properly account
# for the variables local to these generators. Example: `x` is local to
# the expression `z for x in y for z in x`.
node.generators = self.visit_block(node.generators)
if is_dict_comp:
node.key = self.visit(node.key)
node.value = self.visit(node.value)
else:
node.elt = self.visit(node.elt)
return node
def visit_comprehension(self, node):
# It is important to visit children in this order so that the reads to
# the target name are appropriately ignored.
node.iter = self.visit(node.iter)
node.target = self.visit(node.target)
return self.generic_visit(node)
def visit_DictComp(self, node):
return self._process_comprehension(node, is_dict_comp=True)
def visit_ListComp(self, node):
return self._process_comprehension(node, is_list_comp=True)
def visit_SetComp(self, node):
return self._process_comprehension(node)
def visit_GeneratorExp(self, node):
return self._process_comprehension(node)
def visit_ClassDef(self, node):
with self.state[_FunctionOrClass] as fn:
fn.node = node
# The ClassDef node itself has a Scope object that tracks the creation
# of its name, along with the usage of any decorator accompanying it.
self._enter_scope(False)
node.decorator_list = self.visit_block(node.decorator_list)
self.scope.modified.add(qual_names.QN(node.name))
self.scope.bound.add(qual_names.QN(node.name))
node.bases = self.visit_block(node.bases)
node.keywords = self.visit_block(node.keywords)
self._exit_and_record_scope(node)
# A separate Scope tracks the actual class definition.
self._enter_scope(True)
node = self.generic_visit(node)
self._exit_scope()
return node
def _visit_node_list(self, nodes):
return [(None if n is None else self.visit(n)) for n in nodes]
def _visit_arg_annotations(self, node):
node.args.kw_defaults = self._visit_node_list(node.args.kw_defaults)
node.args.defaults = self._visit_node_list(node.args.defaults)
self._track_annotations_only = True
node = self._visit_arg_declarations(node)
self._track_annotations_only = False
return node
def _visit_arg_declarations(self, node):
node.args.posonlyargs = self._visit_node_list(node.args.posonlyargs)
node.args.args = self._visit_node_list(node.args.args)
if node.args.vararg is not None:
node.args.vararg = self.visit(node.args.vararg)
node.args.kwonlyargs = self._visit_node_list(node.args.kwonlyargs)
if node.args.kwarg is not None:
node.args.kwarg = self.visit(node.args.kwarg)
return node
def visit_FunctionDef(self, node):
with self.state[_FunctionOrClass] as fn:
fn.node = node
# The FunctionDef node itself has a Scope object that tracks the creation
# of its name, along with the usage of any decorator accompanying it.
self._enter_scope(False)
node.decorator_list = self.visit_block(node.decorator_list)
if node.returns:
node.returns = self._process_annotation(node.returns)
# Argument annotartions (including defaults) affect the defining context.
node = self._visit_arg_annotations(node)
function_name = qual_names.QN(node.name)
self.scope.modified.add(function_name)
self.scope.bound.add(function_name)
self._exit_and_record_scope(node)
# A separate Scope tracks the actual function definition.
self._enter_scope(True, node.name)
# Keep a separate scope for the arguments node, which is used in the CFG.
self._enter_scope(False, node.name)
# Arg declarations only affect the function itself, and have no effect
# in the defining context whatsoever.
node = self._visit_arg_declarations(node)
self._exit_and_record_scope(node.args)
# Track the body separately. This is for compatibility reasons, it may not
# be strictly needed.
self._enter_scope(False, node.name)
node.body = self.visit_block(node.body)
self._exit_and_record_scope(node, NodeAnno.BODY_SCOPE)
self._exit_and_record_scope(node, NodeAnno.ARGS_AND_BODY_SCOPE)
return node
def visit_Lambda(self, node):
# Lambda nodes are treated in roughly the same way as FunctionDef nodes.
with self.state[_FunctionOrClass] as fn:
fn.node = node
# The Lambda node itself has a Scope object that tracks the creation
# of its name, along with the usage of any decorator accompanying it.
self._enter_scope(False)
node = self._visit_arg_annotations(node)
self._exit_and_record_scope(node)
# A separate Scope tracks the actual function definition.
self._enter_scope(True)
# Keep a separate scope for the arguments node, which is used in the CFG.
self._enter_scope(False)
node = self._visit_arg_declarations(node)
self._exit_and_record_scope(node.args)
# Track the body separately. This is for compatibility reasons, it may not
# be strictly needed.
# TODO(mdan): Do remove it, it's confusing.
self._enter_scope(False)
node.body = self.visit(node.body)
# The lambda body can contain nodes of types normally not found as
# statements, and may not have the SCOPE annotation needed by the CFG.
# So we attach one if necessary.
if not anno.hasanno(node.body, anno.Static.SCOPE):
anno.setanno(node.body, anno.Static.SCOPE, self.scope)
self._exit_and_record_scope(node, NodeAnno.BODY_SCOPE)
lambda_scope = self.scope
self._exit_and_record_scope(node, NodeAnno.ARGS_AND_BODY_SCOPE)
# TODO(bhack:) https://github.com/tensorflow/tensorflow/issues/56089
# remove after deprecation
# Exception: lambdas are assumed to be used in the place where
# they are defined. Therefore, their activity is passed on to the
# calling statement.
self.scope.read.update(lambda_scope.read - lambda_scope.bound)
return node
def visit_With(self, node):
self._enter_scope(False)
node = self.generic_visit(node)
self._exit_and_record_scope(node, NodeAnno.BODY_SCOPE)
return node
def visit_withitem(self, node):
return self._process_statement(node)
def visit_If(self, node):
self._enter_scope(False)
node.test = self.visit(node.test)
node_scope = self._exit_and_record_scope(node.test)
anno.setanno(node, NodeAnno.COND_SCOPE, node_scope)
node = self._process_parallel_blocks(node,
((node.body, NodeAnno.BODY_SCOPE),
(node.orelse, NodeAnno.ORELSE_SCOPE)))
return node
def visit_For(self, node):
self._enter_scope(False)
node.target = self.visit(node.target)
node.iter = self.visit(node.iter)
self._exit_and_record_scope(node.iter)
self._enter_scope(False)
self.visit(node.target)
if anno.hasanno(node, anno.Basic.EXTRA_LOOP_TEST):
self._process_statement(anno.getanno(node, anno.Basic.EXTRA_LOOP_TEST))
self._exit_and_record_scope(node, tag=NodeAnno.ITERATE_SCOPE)
node = self._process_parallel_blocks(node,
((node.body, NodeAnno.BODY_SCOPE),
(node.orelse, NodeAnno.ORELSE_SCOPE)))
return node
def visit_While(self, node):
self._enter_scope(False)
node.test = self.visit(node.test)
node_scope = self._exit_and_record_scope(node.test)
anno.setanno(node, NodeAnno.COND_SCOPE, node_scope)
node = self._process_parallel_blocks(node,
((node.body, NodeAnno.BODY_SCOPE),
(node.orelse, NodeAnno.ORELSE_SCOPE)))
return node
def visit_ExceptHandler(self, node):
self._enter_scope(False)
# try/except oddity: as expected, it leaks any names you defined inside the
# except block, but not the name of the exception variable.
if node.name is not None:
self.scope.isolated_names.add(anno.getanno(node.name, anno.Basic.QN))
node = self.generic_visit(node)
self._exit_scope()
return node
def resolve(node, context, parent_scope=None):
return ActivityAnalyzer(context, parent_scope).visit(node)
@@ -0,0 +1,868 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for activity module."""
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import naming
from tensorflow.python.autograph.pyct import parser
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
from tensorflow.python.platform import test
QN = qual_names.QN
NodeAnno = annos.NodeAnno
global_a = 7
global_b = 17
class ScopeTest(test.TestCase):
def assertMissing(self, qn, scope):
self.assertNotIn(qn, scope.read)
self.assertNotIn(qn, scope.modified)
def assertReadOnly(self, qn, scope):
self.assertIn(qn, scope.read)
self.assertNotIn(qn, scope.modified)
def assertWriteOnly(self, qn, scope):
self.assertNotIn(qn, scope.read)
self.assertIn(qn, scope.modified)
def assertReadWrite(self, qn, scope):
self.assertIn(qn, scope.read)
self.assertIn(qn, scope.modified)
def test_copy_from(self):
scope = activity.Scope(None)
scope.modified.add(QN('foo'))
other = activity.Scope(None)
other.copy_from(scope)
self.assertWriteOnly(QN('foo'), other)
scope.modified.add(QN('bar'))
scope.copy_from(other)
self.assertMissing(QN('bar'), scope)
def test_merge_from(self):
scope = activity.Scope(None)
other = activity.Scope(None)
for col in (scope.modified, scope.read, scope.bound, scope.deleted):
col.add(QN('foo'))
for col in (other.modified, other.read, other.bound, other.deleted):
col.add(QN('foo'))
col.add(QN('bar'))
scope.merge_from(other)
self.assertReadWrite(QN('foo'), scope)
self.assertReadWrite(QN('bar'), scope)
self.assertIn(QN('foo'), scope.bound)
self.assertIn(QN('bar'), scope.bound)
self.assertIn(QN('foo'), scope.deleted)
self.assertIn(QN('bar'), scope.deleted)
def test_copy_of(self):
scope = activity.Scope(None)
scope.read.add(QN('foo'))
other = activity.Scope.copy_of(scope)
self.assertReadOnly(QN('foo'), other)
child_scope = activity.Scope(scope)
child_scope.read.add(QN('bar'))
other = activity.Scope.copy_of(child_scope)
self.assertReadOnly(QN('bar'), other)
def test_referenced(self):
scope = activity.Scope(None)
scope.read.add(QN('a'))
child = activity.Scope(scope)
child.read.add(QN('b'))
child2 = activity.Scope(child, isolated=False)
child2.read.add(QN('c'))
child2.finalize()
child.finalize()
scope.finalize()
self.assertIn(QN('c'), child2.referenced)
self.assertIn(QN('b'), child2.referenced)
self.assertIn(QN('a'), child2.referenced)
self.assertIn(QN('c'), child.referenced)
self.assertIn(QN('b'), child.referenced)
self.assertIn(QN('a'), child.referenced)
class ActivityAnalyzerTestBase(test.TestCase):
def _parse_and_analyze(self, test_fn):
# TODO(mdan): Use a custom FunctionTransformer here.
node, source = parser.parse_entity(test_fn, future_features=())
entity_info = transformer.EntityInfo(
name=test_fn.__name__,
source_code=source,
source_file=None,
future_features=(),
namespace={})
node = qual_names.resolve(node)
namer = naming.Namer({})
ctx = transformer.Context(entity_info, namer, None)
node = activity.resolve(node, ctx)
return node, entity_info
def assertSymbolSetsAre(self, expected, actual, name):
expected = set(expected)
actual = set(str(s) for s in actual)
self.assertSetEqual(
expected, actual, 'for symbol set: %s\n'
' Expected: %s\n'
' Got: %s\n'
' Missing: %s\n'
' Extra: %s\n' % (name.upper(), expected, actual,
expected - actual, actual - expected))
def assertScopeIs(self, scope, used, modified):
"""Assert the scope contains specific used, modified & created variables."""
self.assertSymbolSetsAre(used, scope.read, 'read')
self.assertSymbolSetsAre(modified, scope.modified, 'modified')
class ActivityAnalyzerTest(ActivityAnalyzerTestBase):
def test_import(self):
def test_fn():
import a, b.x, y as c, z.u as d # pylint:disable=g-multiple-import,g-import-not-at-top,unused-variable
node, _ = self._parse_and_analyze(test_fn)
scope = anno.getanno(node.body[0], anno.Static.SCOPE)
self.assertScopeIs(scope, (), ('a', 'b', 'c', 'd'))
def test_import_from(self):
def test_fn():
from x import a # pylint:disable=g-import-not-at-top,unused-variable
from y import z as b # pylint:disable=g-import-not-at-top,unused-variable
node, _ = self._parse_and_analyze(test_fn)
scope = anno.getanno(node.body[0], anno.Static.SCOPE)
self.assertScopeIs(scope, (), ('a',))
scope = anno.getanno(node.body[1], anno.Static.SCOPE)
self.assertScopeIs(scope, (), ('b',))
def test_print_statement(self):
def test_fn(a):
b = 0
c = 1
print(a, b)
return c
node, _ = self._parse_and_analyze(test_fn)
print_node = node.body[2]
if isinstance(print_node, gast.Print):
# Python 2
print_args_scope = anno.getanno(print_node, NodeAnno.ARGS_SCOPE)
else:
# Python 3
assert isinstance(print_node, gast.Expr)
# The call node should be the one being annotated.
print_node = print_node.value
print_args_scope = anno.getanno(print_node, NodeAnno.ARGS_SCOPE)
# We basically need to detect which variables are captured by the call
# arguments.
self.assertScopeIs(print_args_scope, ('a', 'b'), ())
def test_call_args(self):
def test_fn(a):
b = 0
c = 1
foo(a, b) # pylint:disable=undefined-variable
return c
node, _ = self._parse_and_analyze(test_fn)
call_node = node.body[2].value
# We basically need to detect which variables are captured by the call
# arguments.
self.assertScopeIs(
anno.getanno(call_node, NodeAnno.ARGS_SCOPE), ('a', 'b'), ())
def test_call_args_attributes(self):
def foo(*_):
pass
def test_fn(a):
a.c = 0
foo(a.b, a.c)
return a.d
node, _ = self._parse_and_analyze(test_fn)
call_node = node.body[1].value
self.assertScopeIs(
anno.getanno(call_node, NodeAnno.ARGS_SCOPE), ('a', 'a.b', 'a.c'), ())
def test_call_args_subscripts(self):
def foo(*_):
pass
def test_fn(a):
b = 1
c = 2
foo(a[0], a[b])
return a[c]
node, _ = self._parse_and_analyze(test_fn)
call_node = node.body[2].value
self.assertScopeIs(
anno.getanno(call_node, NodeAnno.ARGS_SCOPE),
('a', 'a[0]', 'a[b]', 'b'), ())
def test_while(self):
def test_fn(a):
b = a
while b > 0:
c = b
b -= 1
return b, c
node, _ = self._parse_and_analyze(test_fn)
while_node = node.body[1]
self.assertScopeIs(
anno.getanno(while_node, NodeAnno.BODY_SCOPE), ('b',), ('b', 'c'))
self.assertScopeIs(
anno.getanno(while_node, NodeAnno.BODY_SCOPE).parent, ('a', 'b', 'c'),
('b', 'c'))
self.assertScopeIs(
anno.getanno(while_node, NodeAnno.COND_SCOPE), ('b',), ())
def test_for(self):
def test_fn(a):
b = a
for _ in a:
c = b
b -= 1
return b, c
node, _ = self._parse_and_analyze(test_fn)
for_node = node.body[1]
self.assertScopeIs(
anno.getanno(for_node, NodeAnno.ITERATE_SCOPE), (), ('_'))
self.assertScopeIs(
anno.getanno(for_node, NodeAnno.BODY_SCOPE), ('b',), ('b', 'c'))
self.assertScopeIs(
anno.getanno(for_node, NodeAnno.BODY_SCOPE).parent, ('a', 'b', 'c'),
('b', 'c', '_'))
def test_if(self):
def test_fn(x):
if x > 0:
x = -x
y = 2 * x
z = -y
else:
x = 2 * x
y = -x
u = -y
return z, u
node, _ = self._parse_and_analyze(test_fn)
if_node = node.body[0]
self.assertScopeIs(
anno.getanno(if_node, NodeAnno.BODY_SCOPE), ('x', 'y'), ('x', 'y', 'z'))
self.assertScopeIs(
anno.getanno(if_node, NodeAnno.BODY_SCOPE).parent, ('x', 'y', 'z', 'u'),
('x', 'y', 'z', 'u'))
self.assertScopeIs(
anno.getanno(if_node, NodeAnno.ORELSE_SCOPE), ('x', 'y'),
('x', 'y', 'u'))
self.assertScopeIs(
anno.getanno(if_node, NodeAnno.ORELSE_SCOPE).parent,
('x', 'y', 'z', 'u'), ('x', 'y', 'z', 'u'))
def test_if_attributes(self):
def test_fn(a):
if a > 0:
a.b = -a.c
d = 2 * a
else:
a.b = a.c
d = 1
return d
node, _ = self._parse_and_analyze(test_fn)
if_node = node.body[0]
self.assertScopeIs(
anno.getanno(if_node, NodeAnno.BODY_SCOPE), ('a', 'a.c'), ('a.b', 'd'))
self.assertScopeIs(
anno.getanno(if_node, NodeAnno.ORELSE_SCOPE), ('a', 'a.c'),
('a.b', 'd'))
self.assertScopeIs(
anno.getanno(if_node, NodeAnno.BODY_SCOPE).parent, ('a', 'a.c', 'd'),
('a.b', 'd'))
def test_if_subscripts(self):
def test_fn(a, b, c, e):
if a > 0:
a[b] = -a[c]
d = 2 * a
else:
a[0] = e
d = 1
return d
node, _ = self._parse_and_analyze(test_fn)
if_node = node.body[0]
self.assertScopeIs(
anno.getanno(if_node, NodeAnno.BODY_SCOPE), ('a', 'b', 'c', 'a[c]'),
('a[b]', 'd'))
# TODO(mdan): Should subscript writes (a[0] = 1) be considered to read "a"?
self.assertScopeIs(
anno.getanno(if_node, NodeAnno.ORELSE_SCOPE), ('a', 'e'), ('a[0]', 'd'))
self.assertScopeIs(
anno.getanno(if_node, NodeAnno.ORELSE_SCOPE).parent,
('a', 'b', 'c', 'd', 'e', 'a[c]'), ('d', 'a[b]', 'a[0]'))
def test_nested_if(self):
def test_fn(b):
if b > 0:
if b < 5:
a = b
else:
a = b * b
return a
node, _ = self._parse_and_analyze(test_fn)
inner_if_node = node.body[0].body[0]
self.assertScopeIs(
anno.getanno(inner_if_node, NodeAnno.BODY_SCOPE), ('b',), ('a',))
self.assertScopeIs(
anno.getanno(inner_if_node, NodeAnno.ORELSE_SCOPE), ('b',), ('a',))
def test_nested_function(self):
def test_fn(a):
def f(x):
y = x * x
return y
return f(a)
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'f'), ('f',))
fn_def_node = node.body[0]
scope = anno.getanno(fn_def_node, anno.Static.SCOPE)
self.assertScopeIs(scope, (), ('f'))
scope = anno.getanno(fn_def_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('x', 'y'), ('y',))
scope = anno.getanno(fn_def_node, NodeAnno.ARGS_AND_BODY_SCOPE)
self.assertScopeIs(scope, ('x', 'y'), ('y',))
self.assertSymbolSetsAre(('x', 'y'), scope.bound, 'BOUND')
def test_nested_lambda(self):
def test_fn(a):
return lambda x: (x * a)
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('a',), ())
return_node = node.body[0]
scope = anno.getanno(return_node, anno.Static.SCOPE)
self.assertScopeIs(scope, ('a',), ())
lam_def_node = return_node.value
scope = anno.getanno(lam_def_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'x'), ())
scope = anno.getanno(lam_def_node, NodeAnno.ARGS_AND_BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'x'), ())
self.assertSymbolSetsAre(('x',), scope.bound, 'BOUND')
def test_nested_function_arg_defaults(self):
def test_fn(a):
def f(x=a):
y = x * x
return y
return f(a)
node, _ = self._parse_and_analyze(test_fn)
fn_def_node = node.body[0]
self.assertScopeIs(
anno.getanno(fn_def_node, anno.Static.SCOPE), ('a',), ('f',))
scope = anno.getanno(fn_def_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('x', 'y'), ('y',))
scope = anno.getanno(fn_def_node, NodeAnno.ARGS_AND_BODY_SCOPE)
self.assertScopeIs(scope, ('x', 'y'), ('y',))
self.assertSymbolSetsAre(('x', 'y'), scope.bound, 'BOUND')
def test_constructor_attributes(self):
class TestClass(object):
def __init__(self, a):
self.b = a
self.b.c = 1
node, _ = self._parse_and_analyze(TestClass)
init_node = node.body[0]
self.assertScopeIs(
anno.getanno(init_node, NodeAnno.BODY_SCOPE), ('self', 'a', 'self.b'),
('self', 'self.b', 'self.b.c'))
def test_aug_assign_subscripts(self):
def test_fn(a):
a[0] += 1
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
self.assertScopeIs(
anno.getanno(fn_node, NodeAnno.BODY_SCOPE), ('a', 'a[0]'), ('a[0]',))
def test_return_vars_are_read(self):
def test_fn(a, b, c): # pylint: disable=unused-argument
return c
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
self.assertScopeIs(anno.getanno(fn_node, NodeAnno.BODY_SCOPE), ('c',), ())
self.assertScopeIs(
anno.getanno(node.body[0], anno.Static.SCOPE), ('c',), ())
def test_raise_names_are_read(self):
def test_fn(a, b, c): # pylint: disable=unused-argument
raise b
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
self.assertScopeIs(anno.getanno(fn_node, NodeAnno.BODY_SCOPE), ('b',), ())
self.assertScopeIs(
anno.getanno(node.body[0], anno.Static.SCOPE), ('b',), ())
def test_except_exposes_names(self):
def test_fn(a, b, c): # pylint: disable=unused-argument
try:
pass
except: # pylint: disable=bare-except
b = c
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
self.assertScopeIs(
anno.getanno(fn_node, NodeAnno.BODY_SCOPE), ('c',), ('b',))
def test_except_hides_exception_var_name(self):
def test_fn(a, b, c): # pylint: disable=unused-argument
try:
pass
except a as e:
b = e
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
self.assertScopeIs(
anno.getanno(fn_node, NodeAnno.BODY_SCOPE), ('a',), ('b',))
def test_aug_assign(self):
def test_fn(a, b):
a += b
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
self.assertScopeIs(
anno.getanno(fn_node, NodeAnno.BODY_SCOPE), ('a', 'b'), ('a'))
def test_aug_assign_rvalues(self):
a = dict(bar=3)
def foo():
return a
def test_fn(x):
foo()['bar'] += x
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
self.assertScopeIs(
anno.getanno(fn_node, NodeAnno.BODY_SCOPE), ('foo', 'x'), ())
def test_lambda(self):
def test_fn(a, b):
return lambda: (a + b)
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b'), ())
lam_def_node = node.body[0].value
scope = anno.getanno(lam_def_node, anno.Static.SCOPE)
self.assertScopeIs(scope, (), ())
scope = anno.getanno(lam_def_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b'), ())
scope = anno.getanno(lam_def_node, NodeAnno.ARGS_AND_BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b'), ())
self.assertSymbolSetsAre((), scope.bound, 'BOUND')
scope = anno.getanno(lam_def_node.args, anno.Static.SCOPE)
self.assertSymbolSetsAre((), scope.params.keys(), 'lambda params')
def test_lambda_params_args(self):
def test_fn(a, b): # pylint: disable=unused-argument
return lambda a: a + b
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
# Note: `a` in `a + b` is not "read" here because it's hidden by the `a`
# argument.
self.assertScopeIs(scope, ('b',), ())
lam_def_node = node.body[0].value
scope = anno.getanno(lam_def_node, anno.Static.SCOPE)
self.assertScopeIs(scope, (), ())
scope = anno.getanno(lam_def_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b'), ())
scope = anno.getanno(lam_def_node, NodeAnno.ARGS_AND_BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b'), ())
self.assertSymbolSetsAre(('a',), scope.bound, 'BOUND')
scope = anno.getanno(lam_def_node.args, anno.Static.SCOPE)
self.assertSymbolSetsAre(('a',), scope.params.keys(), 'lambda params')
def test_lambda_params_arg_defaults(self):
def test_fn(a, b, c): # pylint: disable=unused-argument
return lambda b=c: a + b
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
# Note: `b` is not "read" here because it's hidden by the argument.
self.assertScopeIs(scope, ('a', 'c'), ())
lam_def_node = node.body[0].value
scope = anno.getanno(lam_def_node, anno.Static.SCOPE)
self.assertScopeIs(scope, ('c',), ())
scope = anno.getanno(lam_def_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b'), ())
scope = anno.getanno(lam_def_node, NodeAnno.ARGS_AND_BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b'), ())
self.assertSymbolSetsAre(('b',), scope.bound, 'BOUND')
scope = anno.getanno(lam_def_node.args, anno.Static.SCOPE)
self.assertSymbolSetsAre(('b',), scope.params.keys(), 'lambda params')
def test_lambda_complex(self):
def test_fn(a, b, c, d, e): # pylint: disable=unused-argument
a = (lambda a, b, c=e: a + b + c)(d, 1, 2) + b
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('d', 'b', 'e'), ('a',))
lam_def_node = node.body[0].value.left.func
scope = anno.getanno(lam_def_node, anno.Static.SCOPE)
self.assertScopeIs(scope, ('e',), ())
scope = anno.getanno(lam_def_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b', 'c'), ())
scope = anno.getanno(lam_def_node, NodeAnno.ARGS_AND_BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b', 'c'), ())
self.assertSymbolSetsAre(('a', 'b', 'c'), scope.bound, 'BOUND')
scope = anno.getanno(lam_def_node.args, anno.Static.SCOPE)
self.assertSymbolSetsAre(
('a', 'b', 'c'), scope.params.keys(), 'lambda params')
def test_lambda_nested(self):
def test_fn(a, b, c, d, e, f): # pylint: disable=unused-argument
a = lambda a, b: d(lambda b=f: a + b + c) # pylint: disable=undefined-variable
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('d', 'c', 'f'), ('a',))
outer_lam_def = node.body[0].value
scope = anno.getanno(outer_lam_def, anno.Static.SCOPE)
self.assertScopeIs(scope, (), ())
scope = anno.getanno(outer_lam_def, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('d', 'f', 'a', 'c'), ())
scope = anno.getanno(outer_lam_def, NodeAnno.ARGS_AND_BODY_SCOPE)
self.assertScopeIs(scope, ('d', 'f', 'a', 'c'), ())
self.assertSymbolSetsAre(('a', 'b'), scope.bound, 'BOUND')
scope = anno.getanno(outer_lam_def.args, anno.Static.SCOPE)
self.assertSymbolSetsAre(('a', 'b'), scope.params.keys(), 'lambda params')
inner_lam_def = outer_lam_def.body.args[0]
scope = anno.getanno(inner_lam_def, anno.Static.SCOPE)
self.assertScopeIs(scope, ('f',), ())
scope = anno.getanno(inner_lam_def, NodeAnno.BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b', 'c'), ())
scope = anno.getanno(inner_lam_def, NodeAnno.ARGS_AND_BODY_SCOPE)
self.assertScopeIs(scope, ('a', 'b', 'c'), ())
self.assertSymbolSetsAre(('b',), scope.bound, 'BOUND')
scope = anno.getanno(inner_lam_def.args, anno.Static.SCOPE)
self.assertSymbolSetsAre(('b',), scope.params.keys(), 'lambda params')
def test_comprehension_targets_are_isolated(self):
def test_fn(a):
b = {c for c in a} # pylint:disable=unused-variable
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('a',), ('b',))
def test_comprehension_targets_are_isolated_list_function_w_generator(self):
def test_fn(a):
b = list(c for c in a) # pylint:disable=unused-variable
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('a', 'list'), ('b',))
def test_list_comprehension_targets_are_sometimes_isolated(self):
def test_fn(a):
b = [c for c in a] # pylint:disable=unused-variable
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('a',), ('b',))
def test_comprehension_targets_are_isolated_in_augassign(self):
def test_fn(a, b):
b += [c for c in a] # pylint:disable=unused-variable
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('a', 'b'), ('b',))
def test_comprehension_generator_order(self):
def test_fn(a, b, c): # pylint:disable=unused-argument
e = {d: (a, b) for (a, b) in c for d in b} # pylint:disable=unused-variable,g-complex-comprehension
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('c',), ('e',))
def test_global_symbol(self):
def test_fn(c):
global global_a
global global_b
global_a = global_b + c
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('global_a', 'global_b', 'c'), ('global_a',))
self.assertSetEqual(body_scope.globals, set(
(QN('global_a'), QN('global_b'))))
global_a_scope = anno.getanno(fn_node.body[0], anno.Static.SCOPE)
self.assertScopeIs(global_a_scope, ('global_a',), ())
def test_nonlocal_symbol(self):
nonlocal_a = 3
nonlocal_b = 13
def test_fn(c):
nonlocal nonlocal_a
nonlocal nonlocal_b
nonlocal_a = nonlocal_b + c
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(
body_scope, ('nonlocal_a', 'nonlocal_b', 'c'), ('nonlocal_a',))
nonlocal_a_scope = anno.getanno(fn_node.body[0], anno.Static.SCOPE)
self.assertScopeIs(nonlocal_a_scope, ('nonlocal_a',), ())
def test_annotated_assign(self):
b = int
def test_fn(c):
a: b = c
return a
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('b', 'c', 'a'), ('a',))
self.assertSymbolSetsAre(('b',), body_scope.annotations, 'annotations')
ann_assign_scope = anno.getanno(fn_node.body[0], anno.Static.SCOPE)
self.assertScopeIs(ann_assign_scope, ('b', 'c'), ('a',))
self.assertSymbolSetsAre(
('b',), ann_assign_scope.annotations, 'annotations')
def test_pure_definition(self):
b = int
def test_fn():
a: b
return a
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('b', 'a'), ('a',))
self.assertSymbolSetsAre(('b',), body_scope.annotations, 'annotations')
ann_assign_scope = anno.getanno(fn_node.body[0], anno.Static.SCOPE)
self.assertScopeIs(ann_assign_scope, ('b',), ('a',))
self.assertSymbolSetsAre(
('b',), ann_assign_scope.annotations, 'annotations')
def test_function_def_annotations(self):
b = int
c = int
def test_fn(a: b) -> c:
return a
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
fn_scope = anno.getanno(fn_node, anno.Static.SCOPE)
self.assertScopeIs(fn_scope, ('b', 'c'), ('test_fn',))
self.assertSymbolSetsAre(('b', 'c'), fn_scope.annotations, 'annotations')
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('a',), ())
self.assertSymbolSetsAre((), body_scope.annotations, 'annotations')
def test_class_definition_basic(self):
def test_fn(a, b):
class C(a(b)):
d = 1
return C
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('a', 'b', 'C'), ('C',))
def test_class_definition_isolates_method_writes(self):
def test_fn(a, b, c):
class C(a(b)):
d = 1
def e(self):
f = c + 1
return f
return C
node, _ = self._parse_and_analyze(test_fn)
fn_node = node
body_scope = anno.getanno(fn_node, NodeAnno.BODY_SCOPE)
self.assertScopeIs(body_scope, ('a', 'b', 'C', 'c'), ('C',))
if __name__ == '__main__':
test.main()
@@ -0,0 +1,55 @@
# Copyright 2017 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.
# ==============================================================================
"""Annotations used by the static analyzer."""
from enum import Enum
# TODO(mdan): Remove.
class NoValue(Enum):
def __repr__(self): # pylint: disable=invalid-repr-returned
return self.name
class NodeAnno(NoValue):
"""Additional annotations used by the static analyzer.
These are in addition to the basic annotations declared in anno.py.
"""
# Symbols
# These flags are boolean.
IS_LOCAL = 'Symbol is local to the function scope being analyzed.'
IS_PARAM = 'Symbol is a parameter to the function being analyzed.'
IS_MODIFIED_SINCE_ENTRY = (
'Symbol has been explicitly replaced in the current function scope.')
# Scopes
# Scopes are represented by objects of type activity.Scope.
ARGS_SCOPE = 'The scope for the argument list of a function call.'
COND_SCOPE = 'The scope for the test node of a conditional statement.'
ITERATE_SCOPE = 'The scope for the iterate assignment of a for loop.'
ARGS_AND_BODY_SCOPE = (
'The scope for the main body of a function or lambda, including its'
' arguments.')
BODY_SCOPE = (
'The scope for the main body of a statement (True branch for if '
'statements, main body for loops).')
ORELSE_SCOPE = (
'The scope for the orelse body of a statement (False branch for if '
'statements, orelse body for loops).')
@@ -0,0 +1,220 @@
# Copyright 2018 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.
# ==============================================================================
"""Live variable analysis.
See https://en.wikipedia.org/wiki/Live_variable_analysis for a definition of
the following idioms: live variable, live in, live out, which are used
throughout this file.
This analysis attaches the following:
* symbols that are live at the exit of control flow statements
* symbols that are live at the entry of control flow statements
Requires activity analysis.
"""
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import cfg
from tensorflow.python.autograph.pyct import transformer
from tensorflow.python.autograph.pyct.static_analysis import annos
class Analyzer(cfg.GraphVisitor):
"""CFG visitor that performs liveness analysis at statement level."""
def __init__(self, graph, include_annotations):
super(Analyzer, self).__init__(graph)
self.include_annotations = include_annotations
def init_state(self, _):
return set()
def lamba_check(self, fn_ast_node):
if isinstance(fn_ast_node, gast.Lambda):
# Exception: lambda functions are assumed to be used only in the
# place where they are defined, and not later.
return True
return False
def visit_node(self, node):
prev_live_in = self.in_[node]
if anno.hasanno(node.ast_node, anno.Static.SCOPE):
node_scope = anno.getanno(node.ast_node, anno.Static.SCOPE)
gen = node_scope.read
if not self.include_annotations:
gen -= node_scope.annotations
# TODO(mdan): verify whether composites' parents need to be added.
# E.g. whether x needs to be added if x.y is live. Theoretically the
# activity analysis should have both so that wouldn't be needed.
kill = node_scope.modified | node_scope.deleted
live_out = set()
for n in node.next:
live_out |= self.in_[n]
live_in = gen | (live_out - kill)
reaching_functions = anno.getanno(
node.ast_node, anno.Static.DEFINED_FNS_IN)
for fn_ast_node in reaching_functions:
if self.lamba_check(fn_ast_node):
continue
fn_scope = anno.getanno(fn_ast_node, annos.NodeAnno.ARGS_AND_BODY_SCOPE)
# Any closure of a reaching function definition is conservatively
# considered live.
live_in |= (fn_scope.read - fn_scope.bound)
else:
assert self.can_ignore(node), (node.ast_node, node)
live_out = set()
for n in node.next:
live_out |= self.in_[n]
live_in = live_out
self.in_[node] = live_in
self.out[node] = live_out
# TODO(mdan): Move this to the superclass?
return prev_live_in != live_in
class TreeAnnotator(transformer.Base):
"""Runs liveness analysis on each of the functions defined in the AST.
If a function defined other local functions, those will have separate CFGs.
However, dataflow analysis needs to tie up these CFGs to properly emulate the
effect of closures. In the case of liveness, the parent function's live
variables must account for the variables that are live at the entry of each
subfunction. For example:
def foo():
# baz is live from here on
def bar():
print(baz)
This analyzer runs liveness analysis on each individual function, accounting
for the effect above.
"""
def __init__(self, source_info, graphs, include_annotations):
super(TreeAnnotator, self).__init__(source_info)
self.include_annotations = include_annotations
self.allow_skips = False
self.graphs = graphs
self.current_analyzer = None
def visit(self, node):
node = super(TreeAnnotator, self).visit(node)
if (self.current_analyzer is not None and
isinstance(node, gast.stmt) and
node in self.current_analyzer.graph.index):
cfg_node = self.current_analyzer.graph.index[node]
anno.setanno(node, anno.Static.LIVE_VARS_IN,
frozenset(self.current_analyzer.in_[cfg_node]))
return node
def _analyze_function(self, node, is_lambda):
parent_analyzer = self.current_analyzer
analyzer = Analyzer(self.graphs[node], self.include_annotations)
analyzer.visit_reverse()
self.current_analyzer = analyzer
node = self.generic_visit(node)
self.current_analyzer = parent_analyzer
return node
def visit_Lambda(self, node):
return self._analyze_function(node, is_lambda=True)
def visit_FunctionDef(self, node):
return self._analyze_function(node, is_lambda=False)
def _block_statement_live_out(self, node):
successors = self.current_analyzer.graph.stmt_next[node]
stmt_live_out = set()
for s in successors:
stmt_live_out.update(self.current_analyzer.in_[s])
anno.setanno(node, anno.Static.LIVE_VARS_OUT, frozenset(stmt_live_out))
return node
def _block_statement_live_in(self, node, entry_node):
if entry_node in self.current_analyzer.graph.index:
cfg_node = self.current_analyzer.graph.index[entry_node]
stmt_live_in = frozenset(self.current_analyzer.in_[cfg_node])
else:
assert anno.hasanno(entry_node, anno.Static.LIVE_VARS_IN), (
'If not matching a CFG node, must be a block statement:'
' {}'.format(entry_node))
stmt_live_in = anno.getanno(entry_node, anno.Static.LIVE_VARS_IN)
anno.setanno(node, anno.Static.LIVE_VARS_IN, stmt_live_in)
return node
def visit_If(self, node):
node = self.generic_visit(node)
node = self._block_statement_live_out(node)
return self._block_statement_live_in(node, node.test)
def visit_For(self, node):
node = self.generic_visit(node)
node = self._block_statement_live_out(node)
return self._block_statement_live_in(node, node.iter)
def visit_While(self, node):
node = self.generic_visit(node)
node = self._block_statement_live_out(node)
return self._block_statement_live_in(node, node.test)
def visit_Try(self, node):
node = self.generic_visit(node)
node = self._block_statement_live_out(node)
return self._block_statement_live_in(node, node.body[0])
def visit_ExceptHandler(self, node):
node = self.generic_visit(node)
node = self._block_statement_live_out(node)
return self._block_statement_live_in(node, node.body[0])
def visit_With(self, node):
node = self.generic_visit(node)
return self._block_statement_live_in(node, node.items[0])
def visit_Expr(self, node):
node = self.generic_visit(node)
cfg_node = self.current_analyzer.graph.index[node]
anno.setanno(node, anno.Static.LIVE_VARS_OUT,
frozenset(self.current_analyzer.out[cfg_node]))
return node
# TODO(mdan): Investigate the possibility of removing include_annotations.
def resolve(node, source_info, graphs, include_annotations=True):
"""Resolves the live symbols at the exit of control flow statements.
Args:
node: ast.AST
source_info: transformer.SourceInfo
graphs: Dict[ast.FunctionDef, cfg.Graph]
include_annotations: Bool, whether type annotations should be included in
the analysis.
Returns:
ast.AST
"""
node = TreeAnnotator(source_info, graphs, include_annotations).visit(node)
return node
@@ -0,0 +1,30 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for liveness module, that only run in Python 3."""
from tensorflow.python.autograph.pyct.static_analysis import annos
from tensorflow.python.autograph.pyct.static_analysis import liveness_test
from tensorflow.python.platform import test
NodeAnno = annos.NodeAnno
class LivenessAnalyzerTest(liveness_test.LivenessAnalyzerTestBase):
"""Tests which can only run in Python 3."""
if __name__ == '__main__':
test.main()
@@ -0,0 +1,575 @@
# Copyright 2018 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.
# ==============================================================================
"""Tests for liveness module."""
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import cfg
from tensorflow.python.autograph.pyct import naming
from tensorflow.python.autograph.pyct import parser
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 liveness
from tensorflow.python.autograph.pyct.static_analysis import reaching_fndefs
from tensorflow.python.platform import test
global_a = 7
global_b = 17
class LivenessAnalyzerTestBase(test.TestCase):
def _parse_and_analyze(self, test_fn):
# TODO(mdan): Use a custom FunctionTransformer here.
node, source = parser.parse_entity(test_fn, future_features=())
entity_info = transformer.EntityInfo(
name=test_fn.__name__,
source_code=source,
source_file=None,
future_features=(),
namespace={})
node = qual_names.resolve(node)
namer = naming.Namer({})
ctx = transformer.Context(entity_info, namer, None)
node = activity.resolve(node, ctx)
graphs = cfg.build(node)
node = reaching_fndefs.resolve(node, ctx, graphs)
node = liveness.resolve(node, ctx, graphs)
return node
def assertHasLiveOut(self, node, expected):
live_out = anno.getanno(node, anno.Static.LIVE_VARS_OUT)
live_out_strs = set(str(v) for v in live_out)
if not expected:
expected = ()
if not isinstance(expected, tuple):
expected = (expected,)
self.assertSetEqual(live_out_strs, set(expected))
def assertHasLiveIn(self, node, expected):
live_in = anno.getanno(node, anno.Static.LIVE_VARS_IN)
live_in_strs = set(str(v) for v in live_in)
if not expected:
expected = ()
if not isinstance(expected, tuple):
expected = (expected,)
self.assertSetEqual(live_in_strs, set(expected))
class LivenessAnalyzerTest(LivenessAnalyzerTestBase):
def test_live_out_try_block(self):
def test_fn(x, a, b, c): # pylint:disable=unused-argument
if a > 0:
try:
pass
except: # pylint:disable=bare-except
pass
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], 'x')
self.assertHasLiveOut(fn_body[0].body[0], 'x')
def test_live_out_if_inside_except(self):
def test_fn(x, a, b, c): # pylint:disable=unused-argument
if a > 0:
try:
pass
except: # pylint:disable=bare-except
if b > 0:
x = b
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], 'x')
self.assertHasLiveOut(fn_body[0].body[0], 'x')
self.assertHasLiveOut(fn_body[0].body[0].handlers[0].body[0], 'x')
def test_live_out_stacked_if(self):
def test_fn(x, a):
if a > 0:
x = 0
if a > 1:
x = 1
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], ('a', 'x'))
self.assertHasLiveOut(fn_body[1], 'x')
def test_live_out_stacked_if_else(self):
def test_fn(x, a):
if a > 0:
x = 0
if a > 1:
x = 1
else:
x = 2
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], 'a')
self.assertHasLiveOut(fn_body[1], 'x')
def test_live_out_for_basic(self):
def test_fn(x, a):
for i in range(a):
x += i
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], 'x')
def test_live_out_for_iterate(self):
def test_fn(x, a):
for i in range(a):
x += i
return x, i # pylint:disable=undefined-loop-variable
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], ('x', 'i'))
def test_live_out_attributes(self):
def test_fn(x, a):
if a > 0:
x.y = 0
return x.y
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], ('x.y', 'x'))
def test_live_out_nested_functions(self):
def test_fn(a, b):
if b:
a = []
def foo():
return a
foo()
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], 'a')
def test_live_out_nested_functions_defined_ahead(self):
def test_fn(a, b):
def foo():
return a
if b:
a = []
return foo
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[1], ('a', 'foo'))
def test_live_out_nested_functions_defined_after(self):
def test_fn(a, b):
if b:
a = []
def foo():
return a
return foo
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], ('a',))
def test_live_out_lambda(self):
def test_fn(a, b):
if b:
a = []
foo = lambda: a
if b:
pass
return foo
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], ('a', 'b'))
#TODO(@bhack): replace this after deprecation
# https://github.com/tensorflow/tensorflow/issues/56089
self.assertHasLiveOut(fn_body[2], ('foo',))
#self.assertHasLiveOut(fn_body[2], ('a', 'foo'))
def test_live_out_nested_functions_hidden_by_argument(self):
def test_fn(b):
def foo(a):
return a
if b:
a = [] # pylint:disable=unused-variable
return foo
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[1], ('foo'))
def test_live_out_nested_functions_isolation(self):
def test_fn(b):
if b:
a = 0 # pylint:disable=unused-variable
def child():
max(a) # pylint:disable=used-before-assignment
a = 1
return a
child()
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], 'max')
def test_live_out_deletion(self):
def test_fn(x, y, a):
for _ in a:
if x:
del y
else:
y = 0
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[0], ())
def test_live_in_pass(self):
def test_fn(x, a, b, c): # pylint:disable=unused-argument
if a > 0:
pass
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'x'))
self.assertHasLiveIn(fn_body[0].body[0], ('x',))
self.assertHasLiveIn(fn_body[1], ('x',))
def test_live_in_raise(self):
def test_fn(x, a, b, c):
if a > 0:
b = b + 1
raise c
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'b', 'c', 'x'))
self.assertHasLiveIn(fn_body[0].body[0], ('b', 'c'))
self.assertHasLiveIn(fn_body[1], ('x',))
def test_live_out_except_variable(self):
def test_fn(x, a):
try:
pass
except a as b:
raise b
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
# Note: 'a' is not live because there is no raise statement inside the
# try, and we discount the possibility of other code in the try block
# raising an error.
self.assertHasLiveIn(fn_body[0], ('b', 'x'))
def test_live_in_return_statement(self):
def test_fn(x, a, b, c): # pylint:disable=unused-argument
if a > 0:
return x
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'x'))
self.assertHasLiveIn(fn_body[0].body[0], ('x',))
self.assertHasLiveIn(fn_body[1], ('x',))
def test_live_in_try_block(self):
def test_fn(x, a, b, c): # pylint:disable=unused-argument
if a > 0:
try:
pass
except: # pylint:disable=bare-except
pass
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'x'))
self.assertHasLiveIn(fn_body[0].body[0], ('x',))
self.assertHasLiveIn(fn_body[1], ('x',))
def test_live_in_try_orelse(self):
def test_fn(x, a, b, c): # pylint:disable=unused-argument
if a > 0:
try:
pass
except: # pylint:disable=bare-except
pass
else:
x = b
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'b', 'x'))
self.assertHasLiveIn(fn_body[0].body[0], ('b', 'x'))
self.assertHasLiveIn(fn_body[1], ('x',))
def test_live_in_if_inside_except(self):
def test_fn(x, a, b, c): # pylint:disable=unused-argument
if a > 0:
try:
pass
except: # pylint:disable=bare-except
if b > 0:
x = b
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'b', 'x'))
self.assertHasLiveIn(fn_body[0].body[0], ('b', 'x'))
self.assertHasLiveIn(fn_body[0].body[0].handlers[0].body[0], ('b', 'x'))
self.assertHasLiveIn(fn_body[1], ('x',))
def test_live_in_stacked_if(self):
def test_fn(x, a, b, c):
if a > 0:
x = b
if c > 1:
x = 0
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'b', 'c', 'x'))
self.assertHasLiveIn(fn_body[1], ('c', 'x'))
def test_live_in_stacked_if_else(self):
def test_fn(x, a, b, c, d):
if a > 1:
x = b
else:
x = c
if d > 0:
x = 0
return x
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'b', 'c', 'd'))
self.assertHasLiveIn(fn_body[1], ('d', 'x'))
def test_live_in_for_basic(self):
def test_fn(x, y, a):
for i in a:
x = i
y += x
z = 0
return y, z
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'y', 'z'))
def test_live_in_for_nested(self):
def test_fn(x, y, a):
for i in a:
for j in i:
x = i
y += x
z = j
return y, z
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'y', 'z'))
def test_live_in_deletion(self):
def test_fn(x, y, a):
for _ in a:
if x:
del y
else:
y = 0
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('a', 'x', 'y'))
def test_live_in_generator_comprehension(self):
def test_fn(y):
if all(x for x in y):
return
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('all', 'y'))
def test_live_in_list_comprehension(self):
def test_fn(y):
if [x for x in y]:
return
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('y',))
def test_live_in_list_comprehension_expression(self):
def test_fn(y, s):
s += foo([x for x in y]) # pylint:disable=undefined-variable
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('y', 'foo', 's'))
def test_live_in_set_comprehension(self):
def test_fn(y):
if {x for x in y}:
return
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('y',))
def test_live_in_dict_comprehension(self):
def test_fn(y):
if {k: v for k, v in y}:
return
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveIn(fn_body[0], ('y',))
def test_global_symbol(self):
def test_fn(c):
global global_a
global global_b
if global_a:
global_b = c
else:
global_b = c
return global_b
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[2], ('global_b',))
self.assertHasLiveIn(fn_body[2], ('global_a', 'c'))
def test_nonlocal_symbol(self):
nonlocal_a = 3
nonlocal_b = 13
def test_fn(c):
nonlocal nonlocal_a
nonlocal nonlocal_b
if nonlocal_a:
nonlocal_b = c
else:
nonlocal_b = c
return nonlocal_b
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasLiveOut(fn_body[2], ('nonlocal_b',))
self.assertHasLiveIn(fn_body[2], ('nonlocal_a', 'c'))
if __name__ == '__main__':
test.main()
@@ -0,0 +1,288 @@
# Copyright 2018 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.
# ==============================================================================
"""Reaching definition analysis.
This analysis attaches a set of a Definition objects to each symbol, one
for each distinct definition that may reach it. The Definition objects are
mutable and may be used by subsequent analyses to further annotate data like
static type and value information.
The analysis also attaches the set of the symbols defined at the entry of
control flow statements.
Requires activity analysis.
"""
import weakref
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import cfg
from tensorflow.python.autograph.pyct import transformer
class Definition(object):
"""Definition objects describe a unique definition of a variable.
Subclasses of this may be used by passing an appropriate factory function to
resolve.
Attributes:
param_of: Optional[ast.AST]
directives: Dict, optional definition annotations
"""
def __init__(self):
self.param_of = None
self.directives = {}
def __repr__(self):
return '%s[%d]' % (self.__class__.__name__, id(self))
class _NodeState(object):
"""Abstraction for the state of the CFG walk for reaching definition analysis.
This is a value type. Only implements the strictly necessary operators.
Attributes:
value: Dict[qual_names.QN, Set[Definition, ...]], the defined symbols and
their possible definitions
"""
def __init__(self, init_from=None):
if init_from:
if isinstance(init_from, _NodeState):
self.value = {
s: set(other_infos) for s, other_infos in init_from.value.items()
}
elif isinstance(init_from, dict):
self.value = {s: set((init_from[s],)) for s in init_from}
else:
assert False, init_from
else:
self.value = {}
def __eq__(self, other):
if frozenset(self.value.keys()) != frozenset(other.value.keys()):
return False
ret = all(self.value[s] == other.value[s] for s in self.value)
return ret
def __ne__(self, other):
return not self.__eq__(other)
def __or__(self, other):
assert isinstance(other, _NodeState)
result = _NodeState(self)
for s, other_infos in other.value.items():
if s in result.value:
result.value[s].update(other_infos)
else:
result.value[s] = set(other_infos)
return result
def __sub__(self, other):
assert isinstance(other, set)
result = _NodeState(self)
for s in other:
result.value.pop(s, None)
return result
def __repr__(self):
return 'NodeState[%s]=%s' % (id(self), repr(self.value))
class Analyzer(cfg.GraphVisitor):
"""CFG visitor that determines reaching definitions at statement level."""
def __init__(self, graph, definition_factory):
self._definition_factory = definition_factory
super(Analyzer, self).__init__(graph)
self.gen_map = {}
def init_state(self, _):
return _NodeState()
def visit_node(self, node):
prev_defs_out = self.out[node]
defs_in = _NodeState()
for n in node.prev:
defs_in |= self.out[n]
if anno.hasanno(node.ast_node, anno.Static.SCOPE):
node_scope = anno.getanno(node.ast_node, anno.Static.SCOPE)
# The definition objects created by each node must be singletons because
# their ids are used in equality checks.
if node not in self.gen_map:
node_symbols = {}
# Every binding operation (assign, nonlocal, global, etc.) counts as a
# definition, with the exception of del, which only deletes without
# creating a new variable.
newly_defined = ((node_scope.bound | node_scope.globals) -
node_scope.deleted)
for s in newly_defined:
def_ = self._definition_factory()
node_symbols[s] = def_
# Every param receives a definition. Params are not necessarily
# considered as "modified".
for s, p in node_scope.params.items():
def_ = self._definition_factory()
def_.param_of = weakref.ref(p)
node_symbols[s] = def_
self.gen_map[node] = _NodeState(node_symbols)
gen = self.gen_map[node]
kill = node_scope.modified | node_scope.deleted
defs_out = gen | (defs_in - kill)
gen = self.gen_map[node]
defs_out = gen | (defs_in - kill)
else:
assert self.can_ignore(node), (node.ast_node, node)
defs_out = defs_in
self.in_[node] = defs_in
self.out[node] = defs_out
return prev_defs_out != defs_out
class TreeAnnotator(transformer.Base):
"""AST visitor that annotates each symbol name with its reaching definitions.
Simultaneously, the visitor runs the dataflow analysis on each function node,
accounting for the effect of closures. For example:
def foo():
bar = 1
def baz():
# bar = 1 reaches here
"""
def __init__(self, source_info, graphs, definition_factory):
super(TreeAnnotator, self).__init__(source_info)
self.allow_skips = False
self.definition_factory = definition_factory
self.graphs = graphs
self.current_analyzer = None
self.current_cfg_node = None
def visit_FunctionDef(self, node):
parent_analyzer = self.current_analyzer
subgraph = self.graphs[node]
analyzer = Analyzer(subgraph, self.definition_factory)
analyzer.visit_forward()
# Recursively process any remaining subfunctions.
self.current_analyzer = analyzer
node.args = self.visit(node.args)
node.body = self.visit_block(node.body)
self.current_analyzer = parent_analyzer
return node
def visit_Name(self, node):
if self.current_analyzer is None:
# Names may appear outside function defs - for example in class
# definitions.
return node
analyzer = self.current_analyzer
cfg_node = self.current_cfg_node
assert cfg_node is not None, ('name node, %s, outside of any statement?'
% node.id)
qn = anno.getanno(node, anno.Basic.QN)
if isinstance(node.ctx, gast.Load):
anno.setanno(node, anno.Static.DEFINITIONS,
tuple(analyzer.in_[cfg_node].value.get(qn, ())))
else:
anno.setanno(node, anno.Static.DEFINITIONS,
tuple(analyzer.out[cfg_node].value.get(qn, ())))
return node
def _aggregate_predecessors_defined_in(self, node):
preds = self.current_analyzer.graph.stmt_prev[node]
node_defined_in = set()
for p in preds:
node_defined_in |= set(self.current_analyzer.out[p].value.keys())
anno.setanno(node, anno.Static.DEFINED_VARS_IN, frozenset(node_defined_in))
def visit_If(self, node):
self._aggregate_predecessors_defined_in(node)
return self.generic_visit(node)
def visit_For(self, node):
self._aggregate_predecessors_defined_in(node)
# Manually accounting for the shortcoming described in
# cfg.AstToCfg.visit_For.
parent = self.current_cfg_node
self.current_cfg_node = self.current_analyzer.graph.index[node.iter]
node.target = self.visit(node.target)
self.current_cfg_node = parent
node.iter = self.visit(node.iter)
node.body = self.visit_block(node.body)
node.orelse = self.visit_block(node.orelse)
return node
def visit_While(self, node):
self._aggregate_predecessors_defined_in(node)
return self.generic_visit(node)
def visit_Try(self, node):
self._aggregate_predecessors_defined_in(node)
return self.generic_visit(node)
def visit_ExceptHandler(self, node):
self._aggregate_predecessors_defined_in(node)
# TODO(mdan): Also track the exception type / name symbols.
node.body = self.visit_block(node.body)
return node
def visit(self, node):
parent = self.current_cfg_node
if (self.current_analyzer is not None and
node in self.current_analyzer.graph.index):
self.current_cfg_node = self.current_analyzer.graph.index[node]
node = super(TreeAnnotator, self).visit(node)
self.current_cfg_node = parent
return node
def resolve(node, source_info, graphs, definition_factory=Definition):
"""Resolves reaching definitions for each symbol.
Args:
node: ast.AST
source_info: transformer.SourceInfo
graphs: Dict[ast.FunctionDef, cfg.Graph]
definition_factory: Callable[[], Definition]
Returns:
ast.AST
"""
visitor = TreeAnnotator(source_info, graphs, definition_factory)
node = visitor.visit(node)
return node
@@ -0,0 +1,92 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for reaching_definitions module, that only run in Python 3."""
from tensorflow.python.autograph.pyct.static_analysis import reaching_definitions_test
from tensorflow.python.platform import test
class ReachingDefinitionsAnalyzerTest(
reaching_definitions_test.ReachingDefinitionsAnalyzerTestBase):
"""Tests which can only run in Python 3."""
def test_nonlocal(self):
a = 3
b = 13
def test_fn():
nonlocal a
nonlocal b
if a:
b = []
return a, b
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[2].test, 1)
self.assertHasDefs(fn_body[2].body[0].targets[0], 1)
self.assertHasDefs(fn_body[3].value.elts[0], 1)
self.assertHasDefs(fn_body[3].value.elts[1], 2)
self.assertSameDef(fn_body[2].test, fn_body[3].value.elts[0])
self.assertHasDefinedIn(fn_body[2], ('a', 'b'))
def test_nonlocal_in_nested_function(self):
a = 3
b = 13
def test_fn():
a = 3
b = 13
def local_fn():
nonlocal a, b
if a:
b = []
return a, b
return local_fn()
node = self._parse_and_analyze(test_fn)
local_body = node.body[2].body
self.assertHasDefs(local_body[1].test, 1)
self.assertHasDefs(local_body[1].body[0].targets[0], 1)
self.assertHasDefs(local_body[2].value.elts[0], 1)
self.assertHasDefs(local_body[2].value.elts[1], 2)
self.assertSameDef(local_body[1].test, local_body[2].value.elts[0])
# Note: the function name is visible inside the function body. But it's
# a closure variable, not a local.
#
# Example:
#
# >>> def f():
# ... print(f)
# >>> g = f
# >>> f = 'something else'
# >>> g()
# something else
#
self.assertHasDefinedIn(local_body[1], ('a', 'b'))
if __name__ == '__main__':
test.main()
@@ -0,0 +1,526 @@
# Copyright 2018 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.
# ==============================================================================
"""Tests for reaching_definitions module."""
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import cfg
from tensorflow.python.autograph.pyct import naming
from tensorflow.python.autograph.pyct import parser
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 reaching_definitions
from tensorflow.python.platform import test
global_a = 7
global_b = 17
class ReachingDefinitionsAnalyzerTestBase(test.TestCase):
def _parse_and_analyze(self, test_fn):
# TODO(mdan): Use a custom FunctionTransformer here.
node, source = parser.parse_entity(test_fn, future_features=())
entity_info = transformer.EntityInfo(
name=test_fn.__name__,
source_code=source,
source_file=None,
future_features=(),
namespace={})
node = qual_names.resolve(node)
namer = naming.Namer({})
ctx = transformer.Context(entity_info, namer, None)
node = activity.resolve(node, ctx)
graphs = cfg.build(node)
node = reaching_definitions.resolve(node, ctx, graphs,
reaching_definitions.Definition)
return node
def assertHasDefs(self, node, num):
defs = anno.getanno(node, anno.Static.DEFINITIONS)
self.assertEqual(len(defs), num)
for r in defs:
self.assertIsInstance(r, reaching_definitions.Definition)
def assertHasDefinedIn(self, node, expected):
defined_in = anno.getanno(node, anno.Static.DEFINED_VARS_IN)
defined_in_str = set(str(v) for v in defined_in)
if not expected:
expected = ()
if not isinstance(expected, tuple):
expected = (expected,)
self.assertSetEqual(defined_in_str, set(expected))
def assertSameDef(self, first, second):
self.assertHasDefs(first, 1)
self.assertHasDefs(second, 1)
self.assertIs(
anno.getanno(first, anno.Static.DEFINITIONS)[0],
anno.getanno(second, anno.Static.DEFINITIONS)[0])
def assertNotSameDef(self, first, second):
self.assertHasDefs(first, 1)
self.assertHasDefs(second, 1)
self.assertIsNot(
anno.getanno(first, anno.Static.DEFINITIONS)[0],
anno.getanno(second, anno.Static.DEFINITIONS)[0])
class ReachingDefinitionsAnalyzerTest(ReachingDefinitionsAnalyzerTestBase):
def test_conditional(self):
def test_fn(a, b):
a = []
if b:
a = []
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[0].targets[0], 1)
self.assertHasDefs(fn_body[1].test, 1)
self.assertHasDefs(fn_body[1].body[0].targets[0], 1)
self.assertHasDefs(fn_body[2].value, 2)
self.assertHasDefinedIn(fn_body[1], ('a', 'b'))
def test_try_in_conditional(self):
def test_fn(a, b): # pylint:disable=unused-argument
a = []
if b:
try:
pass
except: # pylint:disable=bare-except
pass
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefinedIn(fn_body[1], ('a', 'b'))
self.assertHasDefinedIn(fn_body[1].body[0], ('a', 'b'))
def test_conditional_in_try_in_conditional(self):
def test_fn(a, b):
a = []
if b:
try:
if b:
a = []
except TestException: # pylint:disable=undefined-variable,unused-variable
pass
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefinedIn(fn_body[1], ('a', 'b'))
self.assertHasDefinedIn(fn_body[1].body[0], ('a', 'b'))
# Note: `TestException` and `e` are not tracked.
self.assertHasDefinedIn(fn_body[1].body[0].body[0], ('a', 'b'))
def test_conditional_in_except_in_conditional(self):
def test_fn(a, b):
a = []
if b:
try:
pass
except TestException as e: # pylint:disable=undefined-variable,unused-variable
if b:
a = []
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefinedIn(fn_body[1], ('a', 'b'))
self.assertHasDefinedIn(fn_body[1].body[0], ('a', 'b'))
# Note: `TestException` and `e` are not tracked.
self.assertHasDefinedIn(fn_body[1].body[0].handlers[0].body[0], ('a', 'b'))
def test_while(self):
def test_fn(a):
max(a)
while True:
a = a
a = a
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[0].value.args[0], 1)
self.assertHasDefs(fn_body[1].body[0].targets[0], 1)
self.assertHasDefs(fn_body[1].body[1].targets[0], 1)
self.assertHasDefs(fn_body[1].body[1].value, 1)
# The loop does have an invariant test, but the CFG doesn't know that.
self.assertHasDefs(fn_body[1].body[0].value, 2)
self.assertHasDefs(fn_body[2].value, 2)
def test_while_else(self):
def test_fn(x, i):
y = 0
while x:
x += i
if i:
break
else:
y = 1
return x, y
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[0].targets[0], 1)
self.assertHasDefs(fn_body[1].test, 2)
self.assertHasDefs(fn_body[1].body[0].target, 1)
self.assertHasDefs(fn_body[1].body[1].test, 1)
self.assertHasDefs(fn_body[1].orelse[0].targets[0], 1)
self.assertHasDefs(fn_body[2].value.elts[0], 2)
self.assertHasDefs(fn_body[2].value.elts[1], 2)
def test_for_else(self):
def test_fn(x, i):
y = 0
for i in x:
x += i
if i:
break
else:
continue
else:
y = 1
return x, y
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[0].targets[0], 1)
self.assertHasDefs(fn_body[1].target, 1)
self.assertHasDefs(fn_body[1].body[0].target, 1)
self.assertHasDefs(fn_body[1].body[1].test, 1)
self.assertHasDefs(fn_body[1].orelse[0].targets[0], 1)
self.assertHasDefs(fn_body[2].value.elts[0], 2)
self.assertHasDefs(fn_body[2].value.elts[1], 2)
def test_nested_functions(self):
def test_fn(a, b):
a = []
if b:
a = []
def foo():
return a
foo()
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
def_of_a_in_if = fn_body[1].body[0].targets[0]
self.assertHasDefs(fn_body[0].targets[0], 1)
self.assertHasDefs(fn_body[1].test, 1)
self.assertHasDefs(def_of_a_in_if, 1)
self.assertHasDefs(fn_body[2].value, 2)
inner_fn_body = fn_body[1].body[1].body
def_of_a_in_foo = inner_fn_body[0].value
# Even though `a` is visible in the inner function above, the late binding
# makes it impossible to assume that the same value will be visible at
# call time.
self.assertHasDefs(def_of_a_in_foo, 0)
def test_nested_functions_isolation(self):
def test_fn(a):
a = 0
def child():
a = 1
return a
child()
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
parent_return = fn_body[3]
child_return = fn_body[1].body[1]
# The assignment `a = 1` makes `a` local to `child`.
self.assertNotSameDef(parent_return.value, child_return.value)
def test_function_call_in_with(self):
def foo(_):
pass
def test_fn(a):
with foo(a):
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[0].items[0].context_expr.func, 0)
self.assertHasDefs(fn_body[0].items[0].context_expr.args[0], 1)
def test_mutation_subscript(self):
def test_fn(a):
l = []
l[0] = a
return l
node = self._parse_and_analyze(test_fn)
fn_body = node.body
creation = fn_body[0].targets[0]
mutation = fn_body[1].targets[0].value
use = fn_body[2].value
self.assertSameDef(creation, mutation)
self.assertSameDef(creation, use)
def test_deletion_partial(self):
def test_fn(a):
a = 0
if a:
del a
else:
a = 1
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
first_def = fn_body[0].targets[0]
second_def = fn_body[1].orelse[0].targets[0]
use = fn_body[2].value
self.assertNotSameDef(use, first_def)
self.assertSameDef(use, second_def)
def test_deletion_total(self):
def test_fn(a):
if a:
a = 0
else:
a = 1
del a
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
use = fn_body[2].value
self.assertHasDefs(use, 0)
def test_replacement(self):
def foo(a):
return a
def test_fn(a):
a = foo(a)
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
param = node.args.args[0]
source = fn_body[0].value.args[0]
target = fn_body[0].targets[0]
retval = fn_body[1].value
self.assertSameDef(param, source)
self.assertNotSameDef(source, target)
self.assertSameDef(target, retval)
def test_comprehension_leaking(self):
def test_fn(a):
_ = [x for x in a]
return x # pylint:disable=undefined-loop-variable
node = self._parse_and_analyze(test_fn)
fn_body = node.body
listcomp_target = fn_body[0].value.generators[0].target
retval = fn_body[1].value
# Python2 leaks list comprehension symbols. Python3 doesn't.
# For details, see:
# https://stackoverflow.com/questions/4198906/list-comprehension-rebinds-names-even-after-scope-of-comprehension-is-this-righ
self.assertHasDefs(retval, 0)
def test_function_definition(self):
def test_fn():
def a():
pass
if a: # pylint:disable=using-constant-test
a = None
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[1].test, 1)
self.assertHasDefs(fn_body[1].body[0].targets[0], 1)
self.assertHasDefs(fn_body[2].value, 2)
self.assertHasDefinedIn(fn_body[1], ('a',))
def test_definitions_in_except_block(self):
def test_fn():
try:
pass
except ValueError:
a = None
if a: # pylint:disable=using-constant-test
a = None
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[1].test, 1)
self.assertHasDefs(fn_body[1].body[0].targets[0], 1)
self.assertHasDefs(fn_body[2].value, 2)
self.assertHasDefinedIn(fn_body[1], ('a',))
def test_definitions_in_except_block_of_raising_try(self):
def test_fn():
try:
raise ValueError()
except ValueError:
a = None
if a: # pylint:disable=using-constant-test
a = None
return a
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[1].test, 1)
self.assertHasDefs(fn_body[1].body[0].targets[0], 1)
self.assertHasDefs(fn_body[2].value, 2)
self.assertHasDefinedIn(fn_body[1], ('a',))
def test_global(self):
def test_fn():
global global_a
global global_b
if global_a:
global_b = []
return global_a, global_b
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[2].test, 1)
self.assertHasDefs(fn_body[2].body[0].targets[0], 1)
self.assertHasDefs(fn_body[3].value.elts[0], 1)
self.assertHasDefs(fn_body[3].value.elts[1], 2)
self.assertSameDef(fn_body[2].test, fn_body[3].value.elts[0])
self.assertHasDefinedIn(fn_body[2], ('global_a', 'global_b'))
def test_nonlocal(self):
a = 3
b = 13
def test_fn():
nonlocal a
nonlocal b
if a:
b = []
return a, b
node = self._parse_and_analyze(test_fn)
fn_body = node.body
self.assertHasDefs(fn_body[2].test, 1)
self.assertHasDefs(fn_body[2].body[0].targets[0], 1)
self.assertHasDefs(fn_body[3].value.elts[0], 1)
self.assertHasDefs(fn_body[3].value.elts[1], 2)
self.assertSameDef(fn_body[2].test, fn_body[3].value.elts[0])
self.assertHasDefinedIn(fn_body[2], ('a', 'b'))
def test_nonlocal_in_nested_function(self):
a = 3
b = 13
def test_fn():
a = 3
b = 13
def local_fn():
nonlocal a, b
if a:
b = []
return a, b
return local_fn()
node = self._parse_and_analyze(test_fn)
local_body = node.body[2].body
self.assertHasDefs(local_body[1].test, 1)
self.assertHasDefs(local_body[1].body[0].targets[0], 1)
self.assertHasDefs(local_body[2].value.elts[0], 1)
self.assertHasDefs(local_body[2].value.elts[1], 2)
self.assertSameDef(local_body[1].test, local_body[2].value.elts[0])
# Note: the function name is visible inside the function body. But it's
# a closure variable, not a local.
#
# Example:
#
# >>> def f():
# ... print(f)
# >>> g = f
# >>> f = 'something else'
# >>> g()
# something else
#
self.assertHasDefinedIn(local_body[1], ('a', 'b'))
if __name__ == '__main__':
test.main()
@@ -0,0 +1,178 @@
# Copyright 2018 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.
# ==============================================================================
"""An analysis that determines the reach of a function definition.
A function definition is said to reach a statement if that function may exist
(and therefore may be called) when that statement executes.
"""
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import cfg
from tensorflow.python.autograph.pyct import transformer
class Definition(object):
"""Definition objects describe a unique definition of a function."""
def __init__(self, def_node):
self.def_node = def_node
class _NodeState(object):
"""Abstraction for the state of the CFG walk for reaching definition analysis.
This is a value type. Only implements the strictly necessary operators.
Attributes:
value: Dict[qual_names.QN, Set[Definition, ...]], the defined symbols and
their possible definitions
"""
def __init__(self, init_from=None):
if init_from:
self.value = set(init_from)
else:
self.value = set()
def __eq__(self, other):
return self.value == other.value
def __ne__(self, other):
return self.value != other.value
def __or__(self, other):
assert isinstance(other, _NodeState)
result = _NodeState(self.value)
result.value.update(other.value)
return result
def __add__(self, value):
result = _NodeState(self.value)
result.value.add(value)
return result
def __repr__(self):
return 'NodeState[%s]=%s' % (id(self), repr(self.value))
class Analyzer(cfg.GraphVisitor):
"""CFG visitor that determines reaching definitions at statement level."""
def __init__(self, graph, external_defs):
super(Analyzer, self).__init__(graph)
# This allows communicating that nodes have extra reaching definitions,
# e.g. those that a function closes over.
self.external_defs = external_defs
def init_state(self, _):
return _NodeState()
def visit_node(self, node):
prev_defs_out = self.out[node]
if node is self.graph.entry:
defs_in = _NodeState(self.external_defs)
else:
defs_in = prev_defs_out
for n in node.prev:
defs_in |= self.out[n]
defs_out = defs_in
if isinstance(node.ast_node, (gast.Lambda, gast.FunctionDef)):
defs_out += node.ast_node
self.in_[node] = defs_in
self.out[node] = defs_out
return prev_defs_out != defs_out
class TreeAnnotator(transformer.Base):
"""AST visitor that annotates each symbol name with its reaching definitions.
Simultaneously, the visitor runs the dataflow analysis on each function node,
accounting for the effect of closures. For example:
def foo():
def f():
pass
def g():
# `def f` reaches here
"""
def __init__(self, source_info, graphs):
super(TreeAnnotator, self).__init__(source_info)
self.graphs = graphs
self.allow_skips = False
self.current_analyzer = None
def _proces_function(self, node):
parent_analyzer = self.current_analyzer
subgraph = self.graphs[node]
if (self.current_analyzer is not None
and node in self.current_analyzer.graph.index):
cfg_node = self.current_analyzer.graph.index[node]
defined_in = self.current_analyzer.in_[cfg_node].value
else:
defined_in = ()
analyzer = Analyzer(subgraph, defined_in)
analyzer.visit_forward()
self.current_analyzer = analyzer
node = self.generic_visit(node)
self.current_analyzer = parent_analyzer
return node
def visit_FunctionDef(self, node):
return self._proces_function(node)
def visit_Lambda(self, node):
return self._proces_function(node)
def visit(self, node):
# This can happen before entering the top level function
if (self.current_analyzer is not None
and node in self.current_analyzer.graph.index):
cfg_node = self.current_analyzer.graph.index[node]
anno.setanno(node, anno.Static.DEFINED_FNS_IN,
self.current_analyzer.in_[cfg_node].value)
extra_node = anno.getanno(node, anno.Basic.EXTRA_LOOP_TEST, default=None)
if extra_node is not None:
cfg_node = self.current_analyzer.graph.index[extra_node]
anno.setanno(extra_node, anno.Static.DEFINED_FNS_IN,
self.current_analyzer.in_[cfg_node].value)
return super(TreeAnnotator, self).visit(node)
def resolve(node, source_info, graphs):
"""Resolves reaching definitions for each symbol.
Args:
node: ast.AST
source_info: transformer.SourceInfo
graphs: Dict[ast.FunctionDef, cfg.Graph]
Returns:
ast.AST
"""
visitor = TreeAnnotator(source_info, graphs)
node = visitor.visit(node)
return node
@@ -0,0 +1,54 @@
# Copyright 2018 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.
# ==============================================================================
"""Tests for reaching_fndefs module."""
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import cfg
from tensorflow.python.autograph.pyct import naming
from tensorflow.python.autograph.pyct import parser
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 reaching_definitions
from tensorflow.python.autograph.pyct.static_analysis import reaching_fndefs
from tensorflow.python.platform import test
class ReachingFndefsAnalyzerTest(test.TestCase):
def _parse_and_analyze(self, test_fn):
# TODO(mdan): Use a custom FunctionTransformer here.
node, source = parser.parse_entity(test_fn, future_features=())
entity_info = transformer.EntityInfo(
name=test_fn.__name__,
source_code=source,
source_file=None,
future_features=(),
namespace={})
node = qual_names.resolve(node)
namer = naming.Namer({})
ctx = transformer.Context(entity_info, namer, None)
node = activity.resolve(node, ctx)
graphs = cfg.build(node)
node = reaching_definitions.resolve(node, ctx, graphs)
node = reaching_fndefs.resolve(node, ctx, graphs)
return node
def assertHasFnDefs(self, node):
anno.getanno(node, anno.Static.DEFINED_FNS_IN)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,625 @@
# 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
inference. 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 undocumented - 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
@@ -0,0 +1,942 @@
# 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.
# ==============================================================================
"""Tests for type_inference module."""
from typing import Any, Callable, List
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 transpiler
from tensorflow.python.autograph.pyct.static_analysis import activity
from tensorflow.python.autograph.pyct.static_analysis import reaching_definitions
from tensorflow.python.autograph.pyct.static_analysis import reaching_fndefs
from tensorflow.python.autograph.pyct.static_analysis import type_inference
from tensorflow.python.platform import test
class BasicTestResolver(type_inference.Resolver):
"""A very basic resolver for testing."""
def res_name(self, ns, types_ns, name):
str_name = str(name)
if str_name == 'int':
return {int}, int
return {type(ns[str_name])}, ns[str_name]
def res_value(self, ns, value):
return {type(value)}
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
if type_anno is None:
return None
return {str(type_anno)}
class TestTranspiler(transpiler.GenericTranspiler):
def __init__(self, resolver_type):
super().__init__()
self.resolver = resolver_type()
def get_transformed_name(self, _):
return 'test_item'
def transform_ast(self, node, ctx):
node = qual_names.resolve(node)
node = activity.resolve(node, ctx)
graphs = cfg.build(node)
node = reaching_definitions.resolve(node, ctx, graphs)
node = reaching_fndefs.resolve(node, ctx, graphs)
node = type_inference.resolve(node, ctx, graphs, self.resolver)
return node
class TypeInferenceAnalyzerTest(test.TestCase):
def assertTypes(self, node, expected):
if not isinstance(expected, tuple):
expected = expected,
self.assertSetEqual(
set(anno.getanno(node, anno.Static.TYPES)), set(expected))
def assertClosureTypes(self, node, expected):
actual = anno.getanno(node, anno.Static.CLOSURE_TYPES)
actual = {str(k): v for k, v in actual.items()}
for k, v in expected.items():
self.assertIn(k, actual)
self.assertEqual(actual[k], v)
def test_no_inference_on_unknown_operand_types(self):
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return None
def test_fn(a, b):
return a < b, a - b
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
# With no information on operand types, the operators will infer nothing.
self.assertFalse(
anno.hasanno(fn_body[0].value.elts[0], anno.Static.TYPES))
self.assertFalse(
anno.hasanno(fn_body[0].value.elts[1], anno.Static.TYPES))
def test_resolver_output_checked(self):
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return 1
def test_fn(a):
del a
pass
with self.assertRaisesRegex(ValueError, 'expected to return set'):
TestTranspiler(Resolver).transform(test_fn, None)
def test_argument(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
test_self.assertFalse(f_is_local)
if name == qual_names.QN('a'):
test_self.assertEqual(type_anno, qual_names.QN('int'))
return {str(name) + '_type'}
def test_fn(a: int, b):
return a, b
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value.elts[0], 'a_type')
self.assertTypes(fn_body[0].value.elts[1], 'b_type')
def test_argument_of_local_function(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
if f_name == 'test_fn':
test_self.assertFalse(f_is_local)
test_self.assertEqual(name, qual_names.QN('a'))
test_self.assertEqual(type_anno, qual_names.QN('int'))
elif f_name == 'foo':
test_self.assertTrue(f_is_local)
if name == qual_names.QN('x'):
test_self.assertEqual(type_anno, qual_names.QN('float'))
elif name == qual_names.QN('y'):
test_self.assertIsNone(type_anno)
else:
test_self.fail('unexpected argument {} for {}'.format(name, f_name))
else:
test_self.fail('unexpected function name {}'.format(f_name))
return {str(name) + '_type'}
def test_fn(a: int):
def foo(x: float, y):
return x, y
return foo(a, a)
tr = TestTranspiler(Resolver)
node, _ = tr.transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].body[0].value, (('x_type', 'y_type'),))
self.assertTypes(fn_body[0].body[0].value.elts[0], 'x_type')
self.assertTypes(fn_body[0].body[0].value.elts[1], 'y_type')
def test_assign_straightline(self):
def test_fn(a: int, c: float):
b = a
return a, b, c
node, _ = TestTranspiler(BasicTestResolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].targets[0], 'int')
self.assertTypes(fn_body[0].value, 'int')
self.assertTypes(fn_body[1].value.elts[0], 'int')
self.assertTypes(fn_body[1].value.elts[1], 'int')
self.assertTypes(fn_body[1].value.elts[2], 'float')
def test_expr(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_value(self, ns, value):
test_self.assertEqual(value, tc.a)
return {str}
def res_name(self, ns, types_ns, name):
test_self.assertEqual(name, qual_names.QN('tc'))
return {TestClass}, tc
def res_call(self, ns, types_ns, node, f_type, args, keywords):
test_self.assertEqual(f_type, (str,))
return {int}, None
class TestClass:
def a(self):
pass
tc = TestClass()
def test_fn():
tc.a()
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertEqual(
anno.getanno(fn_body[0].value.func, anno.Static.VALUE), tc.a)
self.assertTypes(fn_body[0].value.func, str)
self.assertTypes(fn_body[0].value, int)
self.assertTypes(fn_body[0], int)
def test_assign_overwriting(self):
def test_fn(a: int, b: float):
c = a
c = b
return c
node, _ = TestTranspiler(BasicTestResolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].targets[0], 'int')
self.assertTypes(fn_body[0].value, 'int')
self.assertTypes(fn_body[1].targets[0], 'float')
self.assertTypes(fn_body[1].value, 'float')
def test_dynamic_attribute_of_static_value(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_value(self, ns, value):
test_self.assertEqual(value, tc.a)
return {int}
def res_name(self, ns, types_ns, name):
test_self.assertEqual(name, qual_names.QN('tc'))
return {TestClass}, tc
class TestClass:
def __init__(self):
self.a = 1
tc = TestClass()
def test_fn():
return tc.a
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value.value, TestClass)
self.assertTypes(fn_body[0].value, int)
self.assertIs(anno.getanno(fn_body[0].value.value, anno.Static.VALUE), tc)
self.assertEqual(anno.getanno(fn_body[0].value, anno.Static.VALUE), tc.a)
def test_static_attribute_of_typed_value(self):
test_self = self
class TestClass:
a = 1
tc = TestClass()
class Resolver(type_inference.Resolver):
def res_name(self, ns, types_ns, name):
test_self.assertEqual(name, qual_names.QN('tc'))
return {TestClass}, None
def res_value(self, ns, value):
test_self.assertIs(value, tc.a)
return {str}
def test_fn():
return tc.a
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value.value, TestClass)
self.assertTypes(fn_body[0].value, str) # Resolver is SOT
self.assertFalse(anno.hasanno(fn_body[0].value.value, anno.Static.VALUE))
self.assertEqual(anno.getanno(fn_body[0].value, anno.Static.VALUE), 1)
def test_static_attribute_of_ambiguous_type(self):
test_self = self
class TestClass1:
a = 1
class TestClass2:
a = 2
tc = TestClass1()
class Resolver(type_inference.Resolver):
def res_name(self, ns, types_ns, name):
test_self.assertEqual(name, qual_names.QN('tc'))
return {TestClass1, TestClass2}, None
def res_value(self, ns, value):
test_self.assertIn(value, (1, 2))
return {str}
def test_fn():
return tc.a
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value.value, (TestClass1, TestClass2))
self.assertFalse(anno.hasanno(fn_body[0].value, anno.Static.TYPES))
self.assertFalse(anno.hasanno(fn_body[0].value.value, anno.Static.VALUE))
self.assertFalse(anno.hasanno(fn_body[0].value, anno.Static.VALUE))
def test_property_of_typed_value(self):
test_self = self
class TestClass:
@property
def a(self):
return 1
tc = TestClass()
class Resolver(type_inference.Resolver):
def res_name(self, ns, types_ns, name):
test_self.assertEqual(name, qual_names.QN('tc'))
return {TestClass}, None
def res_value(self, ns, value):
test_self.assertIs(value, TestClass.a)
test_self.assertNotEqual(value, 1) # Can't evaluate property of class.
return {property}
def test_fn():
return tc.a
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value.value, TestClass)
self.assertTypes(fn_body[0].value, property)
self.assertFalse(anno.hasanno(fn_body[0].value.value, anno.Static.VALUE))
self.assertEqual(
anno.getanno(fn_body[0].value, anno.Static.VALUE), TestClass.a)
def test_dynamic_attribute_of_typed_value(self):
test_self = self
class TestClass:
def __init__(self):
self.a = 1
tc = TestClass()
class Resolver(type_inference.Resolver):
def res_name(self, ns, types_ns, name):
test_self.assertEqual(name, qual_names.QN('tc'))
return {TestClass}, None
def test_fn():
return tc.a
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value.value, TestClass)
self.assertFalse(anno.hasanno(fn_body[0].value, anno.Static.TYPES))
self.assertFalse(anno.hasanno(fn_body[0].value.value, anno.Static.VALUE))
self.assertFalse(anno.hasanno(fn_body[0].value, anno.Static.VALUE))
def test_external_value(self):
a = 'foo'
def test_fn():
b = a
return b
node, _ = TestTranspiler(BasicTestResolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].targets[0], str)
self.assertTypes(fn_body[1].value, str)
def test_external_function(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_name(self, ns, types_ns, name):
test_self.assertEqual(name, qual_names.QN('g'))
return {str}, g
def res_call(self, ns, types_ns, node, f_type, args, keywords):
test_self.assertEqual(f_type, (str,))
test_self.assertEqual(
anno.getanno(node.func, anno.Basic.QN), qual_names.QN('g'))
return {float}, None
def g() -> float:
return 1.0
def test_fn():
a = g()
return a
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value.func, str)
self.assertTypes(fn_body[0].targets[0], float)
self.assertTypes(fn_body[1].value, float)
def test_external_function_side_effects(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_name(self, ns, types_ns, name):
test_self.assertEqual(name, qual_names.QN('g'))
return None, g
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return {str(type_anno)}
def res_call(self, ns, types_ns, node, f_type, args, keywords):
test_self.assertIsNone(f_type)
return None, {qual_names.QN('x'): {str}}
def g():
# The resolver will pretend that this function has the following body:
#
# nonlocal x
# x = 'a'
pass
def test_fn(x: int):
y = x
g()
return x, y
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].targets[0], 'int')
self.assertTypes(fn_body[0].value, 'int')
self.assertTypes(fn_body[2].value.elts[0], str)
self.assertTypes(fn_body[2].value.elts[1], 'int')
def test_local_function_closure(self):
def test_fn(x: int):
def foo():
return x
foo()
node, _ = TestTranspiler(BasicTestResolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].body[0].value, 'int')
self.assertClosureTypes(fn_body[0], {'x': {'int'}})
def test_local_function_closure_nested(self):
def test_fn(x: int):
def foo():
def bar():
return x
bar()
foo()
node, _ = TestTranspiler(BasicTestResolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].body[0].body[0].value, 'int')
self.assertClosureTypes(fn_body[0], {'x': {'int'}})
self.assertClosureTypes(fn_body[0].body[0], {'x': {'int'}})
def test_local_function_closure_mutable_var(self):
def test_fn(x: int):
def foo():
nonlocal x
return x
foo()
node, _ = TestTranspiler(BasicTestResolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].body[1].value, 'int')
self.assertClosureTypes(fn_body[0], {'x': {'int'}})
def test_local_function_closure_ignored_for_bound_symbols(self):
def test_fn(x: float): # pylint:disable=unused-argument
def foo():
x = x + 1 # pylint:disable=used-before-assignment
foo()
node, _ = TestTranspiler(BasicTestResolver).transform(test_fn, None)
fn_body = node.body
self.assertFalse(
anno.hasanno(fn_body[0].body[0].value.left, anno.Static.TYPES))
self.assertClosureTypes(fn_body[0], {'x': {'float'}})
def test_local_function_closure_uses_call_site_types(self):
def test_fn(x: int):
def foo():
return x
x = 1.0
foo()
node, _ = TestTranspiler(BasicTestResolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].body[0].value, float)
self.assertTypes(fn_body[1].targets[0], float)
self.assertClosureTypes(fn_body[0], {'x': {float}})
def test_local_function_hides_locals(self):
def test_fn(a: int): # pylint:disable=unused-argument
def local_fn(v):
a = v
return a
local_fn(1)
node, _ = TestTranspiler(BasicTestResolver).transform(test_fn, None)
fn_body = node.body
self.assertFalse(
anno.hasanno(fn_body[0].body[0].targets[0], anno.Static.TYPES))
def test_local_function_type(self):
def test_fn(x: int):
def foo() -> int:
return x
foo()
node, _ = TestTranspiler(BasicTestResolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[1].value.func, Callable[[Any], int])
self.assertTypes(fn_body[1].value, int)
self.assertTypes(fn_body[1], int)
def test_side_effects_on_arg_function_closure(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_name(self, ns, types_ns, name):
test_self.assertEqual(name, qual_names.QN('g'))
return {Callable[[Callable], None]}, g
def res_value(self, ns, value):
test_self.assertEqual(value, 1.0)
return {float}
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return {str(type_anno)}
def res_call(self, ns, types_ns, node, f_type, args, keywords):
test_self.assertEqual(node.func.id, 'g')
test_self.assertEqual(f_type, (Callable[[Callable], None],))
return None, {qual_names.QN('x'): {str}}
def g(foo):
# The resolver will convey that this function has the following body:
#
# nonlocal x
# x = 'a'
# foo()
del foo
pass
def test_fn(x: int): # pylint:disable=unused-argument
def foo():
return x
x = 1.0
g(foo)
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].body[0].value, str)
def test_subscript(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return {list}
def res_value(self, ns, value):
return {int}
def res_slice(self, ns, types_ns, node, value, slice_):
test_self.assertSetEqual(value, {list})
test_self.assertSetEqual(slice_, {int})
return {str}
def test_fn(a):
return a[1]
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value, str)
self.assertTypes(fn_body[0].value.value, list)
self.assertTypes(fn_body[0].value.slice, int)
def test_tuple_unpacking(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return {list}
def res_value(self, ns, value):
return {int}
def res_slice(self, ns, types_ns, node_or_slice, value, slice_):
test_self.assertIn(node_or_slice, (0, 1))
test_self.assertSetEqual(value, {list})
test_self.assertSetEqual(slice_, {int})
if node_or_slice == 0:
return {float}
else:
return {str}
def test_fn(t):
a, b = t
return a, b
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[1].value, ((float, str),))
self.assertTypes(fn_body[1].value.elts[0], float)
self.assertTypes(fn_body[1].value.elts[1], str)
def test_compare(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return {int}
def res_compare(self, ns, types_ns, node, left, right):
test_self.assertSetEqual(left, {int})
test_self.assertListEqual(right, [{int}])
return {bool}
def test_fn(a, b):
return a < b
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value, bool)
self.assertTypes(fn_body[0].value.left, int)
self.assertTypes(fn_body[0].value.comparators[0], int)
def test_binop(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return {list}
def res_binop(self, ns, types_ns, node, left, right):
test_self.assertSetEqual(left, {list})
test_self.assertSetEqual(right, {list})
return {float}
def test_fn(a, b):
return a @ b
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value, float)
self.assertTypes(fn_body[0].value.left, list)
self.assertTypes(fn_body[0].value.right, list)
def test_unop(self):
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return {list}
def res_unop(self, ns, types_ns, node, opnd):
return {float}
def test_fn(a):
return -a
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value, float)
self.assertTypes(fn_body[0].value.operand, list)
def test_tuple_literal(self):
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return {int}
def test_fn(a, b):
return a, b
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value, ((int, int),))
self.assertTypes(fn_body[0].value.elts[0], int)
self.assertTypes(fn_body[0].value.elts[1], int)
def test_list_literal(self):
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return {int}
def res_list_literal(self, ns, elt_types):
all_types = set()
for s in elt_types:
all_types |= s
return {List[t] for t in all_types}
def test_fn(a, b):
return [a, b]
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[0].value, List[int])
self.assertTypes(fn_body[0].value.elts[0], int)
self.assertTypes(fn_body[0].value.elts[1], int)
def test_tuple_unpacking_syntactic(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
if name == qual_names.QN('a'):
return {int}
else:
return {float}
def res_value(self, ns, value):
test_self.assertIn(value, (0, 1))
return int
def res_slice(self, ns, types_ns, node_or_slice, value, slice_):
test_self.assertIn(node_or_slice, (0, 1))
test_self.assertSetEqual(value, {(int, float)})
test_self.assertEqual(slice_, int)
return {t[node_or_slice] for t in value}
def test_fn(a, b):
c, d = a, b
return c, d
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[1].value, ((int, float),))
self.assertTypes(fn_body[1].value.elts[0], int)
self.assertTypes(fn_body[1].value.elts[1], float)
def test_tuple_unpacking_operational(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
return {(int, float)}
def res_value(self, ns, value):
test_self.assertIn(value, (0, 1))
return int
def res_slice(self, ns, types_ns, node_or_slice, value, slice_):
test_self.assertIn(node_or_slice, (0, 1))
test_self.assertSetEqual(value, {(int, float)})
test_self.assertEqual(slice_, int)
return {t[node_or_slice] for t in value}
def test_fn(a):
c, d = a
return c, d
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
self.assertTypes(fn_body[1].value, ((int, float),))
self.assertTypes(fn_body[1].value.elts[0], int)
self.assertTypes(fn_body[1].value.elts[1], float)
def test_list_expansion_syntactic(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
if name == qual_names.QN('a'):
return {int}
else:
return {float}
def res_value(self, ns, value):
test_self.assertIn(value, (0, 1))
return int
def res_slice(self, ns, types_ns, node_or_slice, value, slice_):
test_self.assertIn(node_or_slice, (0, 1))
test_self.assertSetEqual(value, {(int, float)})
test_self.assertEqual(slice_, int)
return {t[node_or_slice] for t in value}
def test_fn(a, b):
[c, d] = a, b
return c, d
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
# TODO(mdan): Whether it's List or Tuple might be open for interpretation.
self.assertTypes(fn_body[1].value, ((int, float),))
self.assertTypes(fn_body[1].value.elts[0], int)
self.assertTypes(fn_body[1].value.elts[1], float)
def test_list_expansion_operational(self):
test_self = self
class Resolver(type_inference.Resolver):
def res_arg(self, ns, types_ns, f_name, name, type_anno, f_is_local):
if name == qual_names.QN('a'):
return {int}
else:
return {float}
def res_value(self, ns, value):
test_self.assertIn(value, (0, 1))
return int
def res_slice(self, ns, types_ns, node_or_slice, value, slice_):
test_self.assertIn(node_or_slice, (0, 1))
test_self.assertSetEqual(value, {(int, float)})
test_self.assertEqual(slice_, int)
return {t[node_or_slice] for t in value}
def test_fn(a, b):
[c, d] = a, b
return c, d
node, _ = TestTranspiler(Resolver).transform(test_fn, None)
fn_body = node.body
# TODO(mdan): Whether it's List or Tuple might be open for interpretation.
self.assertTypes(fn_body[1].value, ((int, float),))
self.assertTypes(fn_body[1].value.elts[0], int)
self.assertTypes(fn_body[1].value.elts[1], float)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,290 @@
# Copyright 2017 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.
# ==============================================================================
"""AST conversion templates.
Adapted from Tangent.
"""
import ast
import textwrap
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import ast_util
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names
class ContextAdjuster(gast.NodeTransformer):
"""Adjusts the ctx field of nodes to ensure consistency.
This transformer can change the ctx fields of a variable, tuple and other
AST elements that allow one, based on whether the element is being read or
written.
"""
def __init__(self, override_value):
self._ctx_override = override_value
def visit(self, node):
original_override = self._ctx_override
node = super(ContextAdjuster, self).visit(node)
if hasattr(node, 'ctx'):
assert node.ctx is not None, 'node {} has ctx unset'.format(node)
self._ctx_override = original_override
return node
def _apply_override(self, node):
if self._ctx_override is not None:
node.ctx = self._ctx_override()
def visit_Attribute(self, node):
self._apply_override(node)
self._ctx_override = gast.Load
node = self.generic_visit(node)
return node
def visit_Tuple(self, node):
self._apply_override(node)
return self.generic_visit(node)
def visit_List(self, node):
self._apply_override(node)
return self.generic_visit(node)
def visit_Name(self, node):
self._apply_override(node)
return self.generic_visit(node)
def visit_Call(self, node):
self._apply_override(node)
# We may be able to override these to Load(), but for now it's simpler
# to just assert that they're set.
self._ctx_override = None
return self.generic_visit(node)
def visit_Dict(self, node):
# We may be able to override these to Load(), but for now it's simpler
# to just assert that they're set.
self._ctx_override = None
return self.generic_visit(node)
def visit_Subscript(self, node):
self._apply_override(node)
self._ctx_override = gast.Load
node.value = self.visit(node.value)
return self.generic_visit(node)
def visit_comprehension(self, node):
# We may be able to override some of these, but for now it's simpler
# to just assert that they're set.
self._ctx_override = None
return self.generic_visit(node)
def visit_Lambda(self, node):
# We may be able to override some of these, but for now it's simpler
# to just assert that they're set.
self._ctx_override = None
return self.generic_visit(node)
class ReplaceTransformer(gast.NodeTransformer):
"""Replace AST nodes."""
def __init__(self, replacements):
"""Create a new ReplaceTransformer.
Args:
replacements: A mapping from placeholder names to (lists of) AST nodes
that these placeholders will be replaced by.
"""
self.replacements = replacements
self.in_replacements = False
self.preserved_annos = {
anno.Basic.DIRECTIVES,
anno.Basic.EXTRA_LOOP_TEST,
anno.Basic.ORIGIN,
anno.Basic.SKIP_PROCESSING,
anno.Static.ORIG_DEFINITIONS,
'function_context_name',
}
def _prepare_replacement(self, replaced, key):
"""Prepares a replacement AST that's safe to swap in for a node.
Args:
replaced: ast.AST, the node being replaced
key: Hashable, the key of the replacement AST
Returns:
ast.AST, the replacement AST
"""
repl = self.replacements[key]
new_nodes = ast_util.copy_clean(repl, preserve_annos=self.preserved_annos)
if isinstance(new_nodes, gast.AST):
new_nodes = [new_nodes]
return new_nodes
def visit_Expr(self, node):
# When replacing a placeholder with an entire statement, the replacement
# must stand on its own and not be wrapped in an Expr.
new_value = self.visit(node.value)
if new_value is node.value:
return node
return new_value
def visit_keyword(self, node):
if node.arg not in self.replacements:
return self.generic_visit(node)
repl = self._prepare_replacement(node, node.arg)
if isinstance(repl, gast.keyword):
return repl
elif (repl and isinstance(repl, (list, tuple)) and
all(isinstance(r, gast.keyword) for r in repl)):
return repl
# TODO(mdan): We may allow replacing with a string as well.
# For example, if one wanted to replace foo with bar in foo=baz, then
# we could allow changing just node arg, so that we end up with bar=baz.
raise ValueError(
'a keyword argument may only be replaced by another keyword or a '
'non-empty list of keywords. Found: {} for keyword {}'.format(
repl, node.arg))
def visit_FunctionDef(self, node):
node = self.generic_visit(node)
if node.name not in self.replacements:
return node
repl = self.replacements[node.name]
if not isinstance(repl, (gast.Name, ast.Name)):
raise ValueError(
'a function name can only be replaced by a Name node. Found: %s' %
repl)
node.name = repl.id
return node
def visit_Attribute(self, node):
node = self.generic_visit(node)
if node.attr not in self.replacements:
return node
repl = self.replacements[node.attr]
if not isinstance(repl, gast.Name):
raise ValueError(
'An attribute can only be replaced by a Name node. Found: %s' % repl)
node.attr = repl.id
return node
def visit_Name(self, node):
if node.id not in self.replacements:
return node
new_nodes = self._prepare_replacement(node, node.id)
if not new_nodes:
return new_nodes
# Preserve the target context.
adjuster = ContextAdjuster(type(node.ctx))
for n in new_nodes:
if hasattr(n, 'ctx'):
adjuster.visit(n)
if len(new_nodes) == 1:
new_nodes, = new_nodes
return new_nodes
def _convert_to_ast(n):
"""Converts from a known data type to AST."""
# Note: When generating AST nodes from strings/QNs in isolation, ctx is
# unknown. ctx must be filled in according to the template being used.
# See ReplaceTransformer.visit_Name.
if isinstance(n, str):
return gast.Name(id=n, ctx=None, annotation=None, type_comment=None)
if isinstance(n, qual_names.QN):
return n.ast()
if isinstance(n, list):
return [_convert_to_ast(e) for e in n]
if isinstance(n, tuple):
return tuple(_convert_to_ast(e) for e in n)
return n
def replace(template, **replacements):
"""Replaces placeholders in a Python template.
AST Name and Tuple nodes always receive the context that inferred from
the template. However, when replacing more complex nodes (that can potentially
contain Name children), then the caller is responsible for setting the
appropriate context.
Args:
template: A string representing Python code. Any symbol name can be used
that appears in the template code can be used as placeholder.
**replacements: A mapping from placeholder names to (lists of) AST nodes
that these placeholders will be replaced by. String values are also
supported as a shorthand for AST Name nodes with the respective ID.
Returns:
An AST node or list of AST nodes with the replacements made. If the
template was a function, a list will be returned. If the template was a
node, the same node will be returned. If the template was a string, an
AST node will be returned (a `Module` node in the case of a multi-line
string, an `Expr` node otherwise).
Raises:
ValueError: if the arguments are incorrect.
"""
if not isinstance(template, str):
raise ValueError('Expected string template, got %s' % type(template))
for k in replacements:
replacements[k] = _convert_to_ast(replacements[k])
template_str = parser.STANDARD_PREAMBLE + textwrap.dedent(template)
nodes = parser.parse(
template_str,
preamble_len=parser.STANDARD_PREAMBLE_LEN,
single_node=False)
results = []
for node in nodes:
node = ReplaceTransformer(replacements).visit(node)
if isinstance(node, (list, tuple)):
results.extend(node)
else:
results.append(node)
results = [qual_names.resolve(r) for r in results]
return results
def replace_as_expression(template, **replacements):
"""Variant of replace that generates expressions, instead of code blocks."""
replacement = replace(template, **replacements)
if len(replacement) != 1:
raise ValueError(
'single expression expected; for more general templates use replace')
node, = replacement
if isinstance(node, gast.Expr):
return node.value
elif isinstance(node, gast.Name):
return node
raise ValueError(
'the template is expected to generate an expression or a name node;'
' instead found %s' % node)
@@ -0,0 +1,338 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for templates module."""
import types
from absl.testing import parameterized
import gast
from tensorflow.python.autograph.pyct import loader
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names as qn
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.platform import test
class _CtxClearer(gast.NodeTransformer):
def visit(self, node):
super(_CtxClearer, self).visit(node)
if hasattr(node, 'ctx'):
node.ctx = None
return node
def _parse_with_unset_ctx(expr_source):
ast_node = parser.parse_expression(expr_source)
_CtxClearer().visit(ast_node)
return ast_node
class _CtxChecker(gast.NodeTransformer):
def __init__(self, test_instance, expected_ctx):
self.at_top_level = True
self.test_instance = test_instance
self.expected_ctx = expected_ctx
def visit(self, node):
if hasattr(node, 'ctx'):
self.test_instance.assertIsInstance(node.ctx, self.expected_ctx)
if self.at_top_level:
self.at_top_level = False
self.expected_ctx = gast.Load
return super(_CtxChecker, self).visit(node)
class TemplatesTest(test.TestCase, parameterized.TestCase):
def assertExpectedCtxSet(self, node, ctx):
"""Assert that node has ctx=ctx at top and ctx=gast.Load everywhere else."""
checker = _CtxChecker(self, ctx)
checker.visit(node)
def test_replace_tuple(self):
template = """
def test_fn(a, c):
return b,
"""
node = templates.replace(template, b=('a', 'c'))[0]
result, _, _ = loader.load_ast(node)
self.assertEqual((2, 3), result.test_fn(2, 3))
def test_replace_variable(self):
template = """
def test_fn(a):
a += 1
a = 2 * a + 1
return b
"""
node = templates.replace(template, a='b')[0]
result, _, _ = loader.load_ast(node)
self.assertEqual(7, result.test_fn(2))
def test_replace_function_name(self):
template = """
def fname(a):
a += 1
a = 2 * a + 1
return a
"""
node = templates.replace(template, fname='test_fn')[0]
result, _, _ = loader.load_ast(node)
self.assertEqual(7, result.test_fn(2))
def test_replace_code_block(self):
template = """
def test_fn(a):
block
return a
"""
class ShouldBeReplaced(object):
pass
node = templates.replace(
template,
block=[
gast.Assign(
targets=[
gast.Name(
'a',
ctx=ShouldBeReplaced,
annotation=None,
type_comment=None)
],
value=gast.BinOp(
gast.Name(
'a',
ctx=ShouldBeReplaced,
annotation=None,
type_comment=None), gast.Add(),
gast.Constant(1, kind=None)),
),
] * 2)[0]
result, _, _ = loader.load_ast(node)
self.assertEqual(3, result.test_fn(1))
def test_replace_attribute(self):
template = """
def test_fn(a):
return a.foo
"""
node = templates.replace(template, foo='b')[0]
result, _, _ = loader.load_ast(node)
mod = types.ModuleType('test')
mod.b = 3
self.assertEqual(3, result.test_fn(mod))
with self.assertRaises(ValueError):
templates.replace(template, foo=1)
def test_replace_attribute_context(self):
template = """
def test_fn(foo):
foo = 0
"""
node = templates.replace(
template,
foo=parser.parse_expression('a.b.c'))[0]
self.assertIsInstance(node.body[0].targets[0].ctx, gast.Store)
self.assertIsInstance(node.body[0].targets[0].value.ctx, gast.Load)
self.assertIsInstance(node.body[0].targets[0].value.value.ctx, gast.Load)
def test_replace_list_context(self):
template = """
def test_fn(foo):
foo = 0
"""
node = templates.replace(template, foo=parser.parse_expression('[a, b]'))[0]
self.assertIsInstance(node.body[0].targets[0].ctx, gast.Store)
self.assertIsInstance(node.body[0].targets[0].elts[0].ctx, gast.Store)
self.assertIsInstance(node.body[0].targets[0].elts[1].ctx, gast.Store)
def test_replace_tuple_context(self):
template = """
def test_fn(foo):
foo = 0
"""
node = templates.replace(template, foo=parser.parse_expression('(a, b)'))[0]
self.assertIsInstance(node.body[0].targets[0].ctx, gast.Store)
self.assertIsInstance(node.body[0].targets[0].elts[0].ctx, gast.Store)
self.assertIsInstance(node.body[0].targets[0].elts[1].ctx, gast.Store)
def test_replace_expression_context(self):
template = """
def test_fn():
foo
"""
node = templates.replace(
template, foo=parser.parse_expression('a + 2 * b / -c'))[0]
self.assertIsInstance(node.body[0].left.ctx, gast.Load)
self.assertIsInstance(node.body[0].right.left.right.ctx, gast.Load)
def test_replace_complex_context(self):
template = """
def test_fn():
foo = 0
"""
node = templates.replace(
template, foo=parser.parse_expression('bar(([a, b],)).baz'))[0]
self.assertIsInstance(node.body[0].targets[0].ctx, gast.Store)
function_call_arg = node.body[0].targets[0].value.args[0]
self.assertIsInstance(function_call_arg.elts[0].ctx, gast.Load)
self.assertIsInstance(function_call_arg.elts[0].elts[0].ctx, gast.Load)
self.assertIsInstance(function_call_arg.elts[0].elts[1].ctx, gast.Load)
def test_replace_index(self):
template = """
def test_fn():
foo = 0
"""
node = templates.replace(
template, foo=parser.parse_expression('foo(a[b]).bar'))[0]
function_call_arg = node.body[0].targets[0].value.args[0]
self.assertIsInstance(function_call_arg.ctx, gast.Load)
self.assertIsInstance(function_call_arg.slice.ctx, gast.Load)
def test_replace_call_keyword(self):
template = """
def test_fn():
def f(a, d, f):
return a + d + f
return f(1, kws=None)
"""
source = parser.parse_expression('f(d=3, f=5)')
node = templates.replace(template, kws=source.keywords)[0]
result, _, _ = loader.load_ast(node)
self.assertEqual(9, result.test_fn())
with self.assertRaises(ValueError):
templates.replace(template, kws=[])
templates.replace(template, kws=1)
def test_replace_name_with_call(self):
template = """
def test_fn():
b = 5
def g(a):
return 3 * a
def f():
return g
return foo
"""
source = parser.parse_expression('f()(b)')
node = templates.replace(template, foo=source)[0]
result, _, _ = loader.load_ast(node)
self.assertEqual(15, result.test_fn())
def test_replace_name_with_dict(self):
template = """
def test_fn():
return foo['bar']
"""
source = parser.parse_expression('{\'bar\': 3}')
node = templates.replace(template, foo=source)[0]
result, _, _ = loader.load_ast(node)
self.assertEqual(3, result.test_fn())
def test_replace_as_expression(self):
template = """
foo(a)
"""
node = templates.replace_as_expression(template, foo='bar', a='baz')
self.assertIsInstance(node, gast.Call)
self.assertEqual(node.func.id, 'bar')
self.assertEqual(node.args[0].id, 'baz')
def test_replace_as_expression_restrictions(self):
template = """
foo(a)
bar(b)
"""
with self.assertRaises(ValueError):
templates.replace_as_expression(template)
def test_function_call_in_list(self):
template = """
foo(bar)
"""
source = parser.parse_expression('[a(b(1))]')
templates.replace_as_expression(template, bar=source)
def test_star_comprehension_in_function_call(self):
template = """
a = foo(func, args)
"""
source = parser.parse_expression('bar(*[i for i in range(j)])')
node = templates.replace(template, func=source.func, args=source.args)
arg_node = node[0].value.args[1].value
self.assertIsInstance(arg_node.generators[0].target.ctx, gast.Store)
self.assertIsInstance(arg_node.elt.ctx, gast.Load)
def test_lambda_in_function_call(self):
template = """
a = foo(arg)
"""
source = parser.parse_expression('[lambda i: i]')
node = templates.replace(template, arg=source)
lambda_arg = node[0].value.args[0].elts[0]
self.assertIsInstance(lambda_arg.args.args[0].ctx, gast.Param)
self.assertIsInstance(lambda_arg.body.ctx, gast.Load)
def test_replace_name_with_subscript(self):
template = """
foo = bar
"""
replacement = qn.QN(qn.QN('dictionary'), subscript=qn.QN('key'))
node = templates.replace(template, foo=replacement)[0].targets[0]
self.assertIsInstance(node.ctx, gast.Store)
self.assertIsInstance(node.value.ctx, gast.Load)
@parameterized.named_parameters([
('mixed_attr_subscript', 'a.b["c"]'),
('mixed_subscript_attr', 'a[b.c]'),
('nested_subscript', 'a[b[c]]'),
('repeated_subscript', 'a[b][c]'),
])
def test_replace_name_mixed_attr_subscript(self, expression_source):
template = 'foo = bar'
replacement = _parse_with_unset_ctx(expression_source)
target_node = templates.replace(template, foo=replacement)[0].targets[0]
self.assertExpectedCtxSet(target_node, gast.Store)
value_node = templates.replace(template, bar=replacement)[0].value
self.assertExpectedCtxSet(value_node, gast.Load)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,54 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
py_library(
name = "codegen",
srcs = [
"codegen.py",
],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/python/autograph/pyct:templates",
"//third_party/py/numpy",
"@pypi//gast",
],
)
py_library(
name = "decorators",
srcs = ["decorators.py"],
strict_deps = True,
visibility = ["//visibility:public"],
)
py_library(
name = "basic_definitions",
srcs = ["basic_definitions.py"],
strict_deps = True,
visibility = ["//visibility:public"],
)
py_test(
name = "codegen_test",
size = "large",
srcs = ["codegen_test.py"],
strict_deps = True,
tags = [
"manual",
"no_windows",
"nomsan",
"notap",
],
deps = [
":codegen",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow/python/platform:client_testlib",
],
)
@@ -0,0 +1,65 @@
# Copyright 2017 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.
# ==============================================================================
"""Module with basic entity definitions for testing."""
def simple_function(x):
"""Docstring."""
return x # comment
def nested_functions(x):
"""Docstring."""
def inner_fn(y):
return y
return inner_fn(x)
def function_with_print():
print('foo')
simple_lambda = lambda: None
class SimpleClass(object):
def simple_method(self):
return self
def method_with_print(self):
print('foo')
def function_with_multiline_call(x):
"""Docstring."""
return range(
x,
x + 1,
)
def basic_decorator(f):
return f
@basic_decorator
@basic_decorator
def decorated_function(x):
if x > 0:
return 1
return 2
@@ -0,0 +1,230 @@
# Copyright 2017 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.
# ==============================================================================
"""Random code generation for testing/fuzzing."""
# pylint: disable=invalid-name
import random
import string
import gast
import numpy as np
from tensorflow.python.autograph.pyct import templates
class NodeSampler(object):
sample_map = None
def sample(self):
nodes, magnitudes = zip(*self.sample_map.items())
return np.random.choice(
nodes, p=np.array(magnitudes, dtype='float32') / np.sum(magnitudes))
class StatementSampler(NodeSampler):
sample_map = dict((
(gast.Assign, 10),
(gast.Print, 1),
(gast.If, 2),
(gast.While, 2),
(gast.For, 0),
))
class ExpressionSampler(NodeSampler):
sample_map = dict((
(gast.UnaryOp, 1),
(gast.BinOp, 8),
(gast.Name, 1),
(gast.Call, 0),
))
class CompareSampler(NodeSampler):
sample_map = dict((
(gast.Eq, 1),
(gast.NotEq, 1),
(gast.Lt, 1),
(gast.LtE, 1),
(gast.Gt, 1),
(gast.GtE, 1),
(gast.Is, 1),
(gast.IsNot, 1),
))
class BinaryOpSampler(NodeSampler):
sample_map = dict((
(gast.Add, 1),
(gast.Sub, 1),
(gast.Mult, 1),
(gast.Div, 1),
(gast.FloorDiv, 1),
(gast.Mod, 1),
(gast.Pow, 1),
))
class UnaryOpSampler(NodeSampler):
sample_map = dict(((gast.USub, 1), (gast.UAdd, 0)))
class NameSampler(NodeSampler):
sample_map = dict((
('new', 1),
('existing', 1),
))
N_CONTROLFLOW_STATEMENTS = 10
N_FUNCTIONDEF_STATEMENTS = 10
class CodeGenerator(object):
"""Generate random syntactically-valid Python ASTs."""
def __init__(self, max_depth=3, depth=0):
self.max_depth = max_depth
self.depth = depth
def generate_statement(self):
"""Generate a statement node, dispatching to the correct class method."""
desired_node = StatementSampler().sample()
self.depth += 1
# Enforce some constraints on generating statements.
# E.g., if statements need at least 3 readable variables.
# If we fail to satisfy our constraints, draw another sample.
if desired_node in (gast.While, gast.For, gast.If):
if self.depth > self.max_depth:
return self.generate_statement()
# Go get the generator method and run it
method = 'generate_' + desired_node.__name__
visitor = getattr(self, method)
node = visitor()
self.depth -= 1
return node
def sample_node_list(self, low, high, generator):
"""Generate a list of statements of random length.
Args:
low: Fewest number of statements to generate.
high: Highest number of statements to generate.
generator: Function to call to generate nodes.
Returns:
A list of statements.
"""
statements = []
for _ in range(np.random.randint(low, high)):
statements.append(generator())
return statements
def generate_Name(self, ctx=gast.Load()):
variable_name = '_' + ''.join(
random.choice(string.ascii_lowercase) for _ in range(4))
return gast.Name(variable_name, ctx=ctx, annotation=None)
def generate_BinOp(self):
# TODO(alexbw): convert to generate_expression when we get to limit
# expression depth.
op = BinaryOpSampler().sample()()
return gast.BinOp(self.generate_Name(), op, self.generate_Name())
def generate_Compare(self):
op = CompareSampler().sample()()
return gast.Compare(self.generate_Name(), [op], [self.generate_Name()])
def generate_UnaryOp(self):
operand = self.generate_Name()
op = UnaryOpSampler().sample()()
return gast.UnaryOp(op, operand)
def generate_expression(self):
desired_node = ExpressionSampler().sample()
# Go get the generator method and run it
method = 'generate_' + desired_node.__name__
generator = getattr(self, method)
return generator()
def generate_Assign(self):
"""Generate an Assign node."""
# Generate left-hand side
target_node = self.generate_Name(gast.Store())
# Generate right-hand side
value_node = self.generate_expression()
# Put it all together
node = gast.Assign(targets=[target_node], value=value_node)
return node
def generate_If(self):
"""Generate an If node."""
test = self.generate_Compare()
# Generate true branch statements
body = self.sample_node_list(
low=1,
high=N_CONTROLFLOW_STATEMENTS // 2,
generator=self.generate_statement)
# Generate false branch statements
orelse = self.sample_node_list(
low=1,
high=N_CONTROLFLOW_STATEMENTS // 2,
generator=self.generate_statement)
node = gast.If(test, body, orelse)
return node
def generate_While(self):
"""Generate a While node."""
test = self.generate_Compare()
body = self.sample_node_list(
low=1, high=N_CONTROLFLOW_STATEMENTS, generator=self.generate_statement)
orelse = [] # not generating else statements
node = gast.While(test, body, orelse)
return node
def generate_Call(self):
raise NotImplementedError
def generate_Return(self):
return gast.Return(self.generate_expression())
def generate_Print(self):
return templates.replace('print(x)', x=self.generate_expression())[0]
def generate_FunctionDef(self):
"""Generate a FunctionDef node."""
# Generate the arguments, register them as available
arg_vars = self.sample_node_list(
low=2, high=10, generator=lambda: self.generate_Name(gast.Param()))
args = gast.arguments(arg_vars, None, [], [], None, [])
# Generate the function body
body = self.sample_node_list(
low=1, high=N_FUNCTIONDEF_STATEMENTS, generator=self.generate_statement)
body.append(self.generate_Return())
fn_name = self.generate_Name().id
node = gast.FunctionDef(fn_name, args, body, (), None)
return node
def generate_random_functiondef():
return CodeGenerator().generate_FunctionDef()
@@ -0,0 +1,36 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for type_info module."""
import numpy as np
from tensorflow.python.autograph.pyct import compiler
from tensorflow.python.autograph.pyct.testing import codegen
from tensorflow.python.platform import test
class CodeGenTest(test.TestCase):
def test_codegen_gens(self):
np.random.seed(0)
for _ in range(1000):
node = codegen.generate_random_functiondef()
fn = compiler.ast_to_object(node)
self.assertIsNotNone(
fn, 'Generated invalid AST that could not convert to source.')
if __name__ == '__main__':
test.main()
@@ -0,0 +1,46 @@
# Copyright 2017 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.
# ==============================================================================
"""Module with test decorators."""
import functools
def wrapping_decorator(f):
@functools.wraps(f)
def wrapper(*args, **kwargs):
return f(*args, **kwargs)
return wrapper
def standalone_decorator(f):
def standalone_wrapper(*args, **kwargs):
return f(*args, **kwargs)
return standalone_wrapper
def functional_decorator():
def decorator(f):
def functional_wrapper(*args, **kwargs):
return f(*args, **kwargs)
return functional_wrapper
return decorator
@@ -0,0 +1,538 @@
# Copyright 2017 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.
# ==============================================================================
"""A node transformer that includes utilities for SCT."""
import collections
import enum
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import pretty_printer
from tensorflow.python.autograph.pyct import templates
class AnalysisLevel(enum.IntEnum):
NONE = 0
ACTIVITY = 1
DEFINEDNESS = 2
LIVENESS = 3
# TODO(znado): Use namedtuple.
class Context(object):
"""Contains information about a source code transformation.
This object is mutable, and is updated during conversion. Not thread safe.
Attributes:
info: EntityInfo, immutable.
namer: naming.Namer.
current_origin: origin_info.OriginInfo, holds the OriginInfo of the last
AST node to be processed successfully. Useful for error handling.
user: An user-supplied context object. The object is opaque to the
infrastructure, but will pe passed through to all custom transformations.
"""
def __init__(self, info, namer, user_context):
self.info = info
self.namer = namer
self.current_origin = None
self.user = user_context
# TODO(mdan): Move to a standalone file.
class EntityInfo(
collections.namedtuple(
'EntityInfo',
('name', 'source_code', 'source_file', 'future_features', 'namespace'))
):
"""Contains information about a Python entity.
Immutable.
Examples of entities include functions and classes.
Attributes:
name: The name that identifies this entity.
source_code: The entity's source code.
source_file: The entity's source file.
future_features: Tuple[Text], the future features that this entity was
compiled with. See
https://docs.python.org/2/reference/simple_stmts.html#future.
namespace: Dict[str, ], containing symbols visible to the entity (excluding
parameters).
"""
pass
class _StateStack(object):
"""Templated context manager.
This class provides syntactic sugar for a stack of objects of known
type. It allows accessing attributes of the object at the top of the stack
directly against this object, which allows for very terse syntax.
For example, this code:
stack = _StateStack(Foo)
stack.enter()
stack.bar
Is equivalent to:
stack = []
stack.append(Foo())
foo = stack[-1]
foo.bar
See _State for more on how this is used.
Attributes:
type: Any, the type of objects that this stack holds
level: int, the current stack depth
stack: List[Any], the actual stack
value: Any, the instance of the object at the top of the stack
"""
def __init__(self, type_):
# Because we override __setattr__, we need to attach these attributes using
# the superclass' setattr.
object.__setattr__(self, 'type', type_)
object.__setattr__(self, '_stack', [])
if not hasattr(type_, 'no_root'):
self.enter()
def __enter__(self):
self.enter()
return self
def __exit__(self, exc_type, exc_value, traceback):
self.exit()
def enter(self):
self._stack.append(self.type())
def exit(self):
self._stack.pop()
@property
def stack(self):
return self._stack
@property
def level(self):
return len(self._stack)
@property
def value(self):
return self._stack[-1]
def __iter__(self):
return iter(self._stack)
def __getattr__(self, key):
return getattr(self._stack[-1], key)
def __setattr__(self, key, value):
setattr(self._stack[-1], key, value)
class _State(object):
"""Syntactic sugar for accessing an instance of a StateStack context manager.
This structure offers syntactic sugar over a dict of stacks of objects
of known type. These structures are useful to keep state during AST walks.
Multiple different scopes can be tracked in parallel. For example:
s = _State()
s[foo].enter()
s[bar].enter() # this will not affect s[foo]
Element access has special semantics:
* keys are a data type
* element values are _StateStack(type=key) objects
* missing elements are automatically added, similarly to defaultdict
For example, the following block :
_State s
s[Foo]
Is equivalent to:
s = {}
if Foo not in s:
s[Foo] = Foo()
s[Foo]
See Base for how it's used.
"""
def __init__(self):
self._value = {}
def __getitem__(self, key):
if key not in self._value:
self._value[key] = _StateStack(key)
return self._value[key]
class NodeStateTracker(object):
"""Base class for general-purpose Python code transformation.
This abstract class provides helpful functions, like state tracking within
the scope of arbitrary node, helpers for processing code blocks, debugging,
mapping of transformed code to original code, and others.
Scope-local state tracking: to keep state across nodes, at the level of
(possibly nested) scopes, use enter/exit_local_scope and set/get_local.
You must call enter/exit_local_scope manually, but the transformer detects
when they are not properly paired.
The transformer allows keeping state across calls that is local
to arbitrary nodes and their descendants, using the self.state attribute.
Multiple independent scopes are allowed and automatically constructed.
For example, to keep track of the `If` node that encloses any `Name` node,
one can write:
```
class FooType(object):
def __init__(self):
self.foo_property = None
class DummyTransformer(NodeStateTracker, ast.NodeTransformer):
def visit_If(self, node):
self.state[FooType].enter()
self.state[FooType].foo_property = node
node = self.veneric_visit(node)
self.state[FooType].exit()
return node
def visit_Name(self, node):
self.state[FooType].foo_property # will hold the innermost enclosing if
```
Alternatively, the `enter()`/`exit()` calls can be managed by a `with`
statement:
```
def visit_If(self, node):
with self.state[FooType] as foo:
foo.foo_property = node
return self.generic_visit(node)
```
"""
# TODO(mdan): Document all extra features.
def __init__(self, ctx):
"""Initialize the transformer.
Subclasses should call this.
Args:
ctx: A Context object.
"""
self._lineno = 0
self._col_offset = 0
self.ctx = ctx
# Allows scoping of local variables to keep state across calls to visit_*
# methods. Multiple scope hierarchies may exist and are keyed by tag. A
# scope is valid at one or more nodes and all its children. Scopes created
# in child nodes supersede their parent. Scopes are isolated from one
# another.
self.state = _State()
def debug_print(self, node):
"""Helper method useful for debugging. Prints the AST."""
if __debug__:
print(pretty_printer.fmt(node))
return node
def debug_print_src(self, node):
"""Helper method useful for debugging. Prints the AST as code."""
if __debug__:
print(parser.unparse(node))
return node
def visit_block(self, nodes, before_visit=None, after_visit=None):
"""A more powerful version of generic_visit for statement blocks.
An example of a block is the body of an if statement.
This function allows specifying a postprocessing callback (the
after_visit argument) argument which can be used to move nodes to a new
destination. This is done by after_visit by returning a non-null
second return value, e.g. return new_node, new_destination.
For example, a transformer could perform the following move:
foo()
bar()
baz()
foo()
if cond:
bar()
baz()
The above could be done with a postprocessor of this kind:
def after_visit(node):
if node_is_function_call(bar):
new_container_node = build_cond()
new_container_node.body.append(node)
return new_container_node, new_container_node.body
else:
# Once we set a new destination, all subsequent items will be
# moved to it, so we don't need to explicitly handle baz.
return node, None
Args:
nodes: enumerable of AST node objects. If None, the function returns None.
before_visit: optional callable that is called before visiting each item
in nodes
after_visit: optional callable that takes in an AST node and returns a
tuple (new_node, new_destination). It is called after visiting each item
in nodes. Is used in the same was as the visit_* methods: new_node will
replace the node; if not None, new_destination must be a list, and
subsequent nodes will be placed in this list instead of the list
returned by visit_block.
Returns:
A list of AST node objects containing the transformed items from nodes,
except those nodes that have been relocated using after_visit.
"""
if nodes is None:
return None
results = []
node_destination = results
for node in nodes:
if before_visit:
# TODO(mdan): We can modify node here too, if ever needed.
before_visit()
replacement = self.visit(node)
if after_visit and replacement:
replacement, new_destination = after_visit(replacement)
else:
new_destination = None
if replacement:
if isinstance(replacement, (list, tuple)):
node_destination.extend(replacement)
else:
node_destination.append(replacement)
# Allow the postprocessor to reroute the remaining nodes to a new list.
if new_destination is not None:
node_destination = new_destination
return results
# TODO(mdan): Rename to PythonCodeTransformer.
class Base(NodeStateTracker, gast.NodeTransformer):
"""Base class for general-purpose Python-to-Python code transformation.
This is an extension of ast.NodeTransformer that provides the additional
functions offered by NodeStateTracker.
"""
def create_assignment(self, target, expression):
template = """
target = expression
"""
return templates.replace(template, target=target, expression=expression)
# TODO(mdan): Remove.
def apply_to_single_assignments(self, targets, values, apply_fn):
"""Applies a function to each individual assignment.
This function can process a possibly-unpacked (e.g. a, b = c, d) assignment.
It tries to break down the unpacking if possible. In effect, it has the same
effect as passing the assigned values in SSA form to apply_fn.
Examples:
The following will result in apply_fn(a, c), apply_fn(b, d):
a, b = c, d
The following will result in apply_fn(a, c[0]), apply_fn(b, c[1]):
a, b = c
The following will result in apply_fn(a, (b, c)):
a = b, c
It uses the visitor pattern to allow subclasses to process single
assignments individually.
Args:
targets: list, tuple of or individual AST node. Should be used with the
targets field of an ast.Assign node.
values: an AST node.
apply_fn: a function of a single argument, which will be called with the
respective nodes of each single assignment. The signature is
apply_fn(target, value), no return value.
"""
if not isinstance(targets, (list, tuple)):
targets = (targets,)
for target in targets:
if isinstance(target, (gast.Tuple, gast.List)):
for i in range(len(target.elts)):
target_el = target.elts[i]
if isinstance(values, (gast.Tuple, gast.List)):
value_el = values.elts[i]
else:
value_el = gast.Subscript(values, i, ctx=gast.Store())
self.apply_to_single_assignments(target_el, value_el, apply_fn)
else:
# TODO(mdan): Look into allowing to rewrite the AST here.
apply_fn(target, values)
def visit(self, node):
if not isinstance(node, gast.AST):
# This is not that uncommon a mistake: various node bodies are lists, for
# example, posing a land mine for transformers that need to recursively
# call `visit`. The error needs to be raised before the exception handler
# below is installed, because said handler will mess up if `node` is not,
# in fact, a node.
msg = ('invalid value for "node": expected "ast.AST", got "{}"; to'
' visit lists of nodes, use "visit_block" instead').format(
type(node))
raise ValueError(msg)
if anno.hasanno(node, anno.Basic.SKIP_PROCESSING):
return node
parent_origin = self.ctx.current_origin
if anno.hasanno(node, anno.Basic.ORIGIN):
self.ctx.current_origin = anno.getanno(node, anno.Basic.ORIGIN)
try:
processing_expr_node = isinstance(node, gast.Expr)
if processing_expr_node:
entry_expr_value = node.value
result = super(Base, self).visit(node)
# Adjust for consistency: replacing the value of an Expr with
# an Assign node removes the need for the Expr node.
if (processing_expr_node and isinstance(result, gast.Expr) and
(result.value is not entry_expr_value)):
# When the replacement is a list, it is assumed that the list came
# from a template that contained a number of statements, which
# themselves are standalone and don't require an enclosing Expr.
if isinstance(result.value,
(list, tuple, gast.Assign, gast.AugAssign)):
result = result.value
# By default, all replacements receive the origin info of the replaced
# node.
if result is not node and result is not None:
inherited_origin = anno.getanno(
node, anno.Basic.ORIGIN, default=parent_origin)
if inherited_origin is not None:
nodes_to_adjust = result
if isinstance(result, (list, tuple)):
nodes_to_adjust = result
else:
nodes_to_adjust = (result,)
for n in nodes_to_adjust:
if not anno.hasanno(n, anno.Basic.ORIGIN):
anno.setanno(n, anno.Basic.ORIGIN, inherited_origin)
finally:
self.ctx.current_origin = parent_origin
return result
class CodeGenerator(NodeStateTracker, gast.NodeVisitor):
"""Base class for general-purpose Python-to-string code transformation.
Similar to Base, but outputs arbitrary strings instead of a Python AST.
This uses the same visitor mechanism that the standard NodeVisitor uses,
meaning that subclasses write handlers for the different kinds of nodes.
New code is generated using the emit method, which appends to a code buffer
that can be afterwards obtained from code_buffer.
Example:
class SimpleCodeGen(CodeGenerator):
def visitIf(self, node):
self.emit('if ')
self.visit(node.test)
self.emit(' { ')
self.visit(node.body)
self.emit(' } else { ')
self.visit(node.orelse)
self.emit(' } ')
node = ast.parse(...)
gen = SimpleCodeGen()
gen.visit(node)
# gen.code_buffer contains the resulting code
"""
def __init__(self, ctx):
super(CodeGenerator, self).__init__(ctx)
self._output_code = ''
self.source_map = {}
def emit(self, code):
self._output_code += code
@property
def code_buffer(self):
return self._output_code
def visit(self, node):
if anno.hasanno(node, anno.Basic.SKIP_PROCESSING):
return
parent_origin = self.ctx.current_origin
eof_before = len(self._output_code)
if anno.hasanno(node, anno.Basic.ORIGIN):
self.ctx.current_origin = anno.getanno(node, anno.Basic.ORIGIN)
try:
ret = super(CodeGenerator, self).visit(node)
# By default, all replacements receive the origin info of the replaced
# node.
eof_after = len(self._output_code)
if eof_before - eof_after:
inherited_origin = anno.getanno(
node, anno.Basic.ORIGIN, default=parent_origin)
if inherited_origin is not None:
self.source_map[(eof_before, eof_after)] = inherited_origin
return ret
finally:
self.ctx.current_origin = parent_origin
@@ -0,0 +1,364 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for templates module."""
import re
import gast
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import origin_info
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import transformer
from tensorflow.python.platform import test
class TransformerTest(test.TestCase):
def _simple_context(self):
entity_info = transformer.EntityInfo(
name='Test_fn',
source_code=None,
source_file=None,
future_features=(),
namespace=None)
return transformer.Context(entity_info, None, None)
def assertSameAnno(self, first, second, key):
self.assertIs(anno.getanno(first, key), anno.getanno(second, key))
def assertDifferentAnno(self, first, second, key):
self.assertIsNot(anno.getanno(first, key), anno.getanno(second, key))
def test_state_tracking(self):
class LoopState(object):
pass
class CondState(object):
pass
class TestTransformer(transformer.Base):
def visit(self, node):
anno.setanno(node, 'loop_state', self.state[LoopState].value)
anno.setanno(node, 'cond_state', self.state[CondState].value)
return super(TestTransformer, self).visit(node)
def visit_While(self, node):
self.state[LoopState].enter()
node = self.generic_visit(node)
self.state[LoopState].exit()
return node
def visit_If(self, node):
self.state[CondState].enter()
node = self.generic_visit(node)
self.state[CondState].exit()
return node
tr = TestTransformer(self._simple_context())
def test_function(a):
a = 1
while a:
_ = 'a'
if a > 2:
_ = 'b'
while True:
raise '1'
if a > 3:
_ = 'c'
while True:
raise '1'
node, _ = parser.parse_entity(test_function, future_features=())
node = tr.visit(node)
fn_body = node.body
outer_while_body = fn_body[1].body
self.assertSameAnno(fn_body[0], outer_while_body[0], 'cond_state')
self.assertDifferentAnno(fn_body[0], outer_while_body[0], 'loop_state')
first_if_body = outer_while_body[1].body
self.assertDifferentAnno(outer_while_body[0], first_if_body[0],
'cond_state')
self.assertSameAnno(outer_while_body[0], first_if_body[0], 'loop_state')
first_inner_while_body = first_if_body[1].body
self.assertSameAnno(first_if_body[0], first_inner_while_body[0],
'cond_state')
self.assertDifferentAnno(first_if_body[0], first_inner_while_body[0],
'loop_state')
second_if_body = outer_while_body[2].body
self.assertDifferentAnno(first_if_body[0], second_if_body[0], 'cond_state')
self.assertSameAnno(first_if_body[0], second_if_body[0], 'loop_state')
second_inner_while_body = second_if_body[1].body
self.assertDifferentAnno(first_inner_while_body[0],
second_inner_while_body[0], 'cond_state')
self.assertDifferentAnno(first_inner_while_body[0],
second_inner_while_body[0], 'loop_state')
def test_state_tracking_context_manager(self):
class CondState(object):
pass
class TestTransformer(transformer.Base):
def visit(self, node):
anno.setanno(node, 'cond_state', self.state[CondState].value)
return super(TestTransformer, self).visit(node)
def visit_If(self, node):
with self.state[CondState]:
return self.generic_visit(node)
tr = TestTransformer(self._simple_context())
def test_function(a):
a = 1
if a > 2:
_ = 'b'
if a < 5:
_ = 'c'
_ = 'd'
node, _ = parser.parse_entity(test_function, future_features=())
node = tr.visit(node)
fn_body = node.body
outer_if_body = fn_body[1].body
self.assertDifferentAnno(fn_body[0], outer_if_body[0], 'cond_state')
self.assertSameAnno(outer_if_body[0], outer_if_body[2], 'cond_state')
inner_if_body = outer_if_body[1].body
self.assertDifferentAnno(inner_if_body[0], outer_if_body[0], 'cond_state')
def test_visit_block_postprocessing(self):
class TestTransformer(transformer.Base):
def _process_body_item(self, node):
if isinstance(node, gast.Assign) and (node.value.id == 'y'):
if_node = gast.If(
gast.Name(
'x', ctx=gast.Load(), annotation=None, type_comment=None),
[node], [])
return if_node, if_node.body
return node, None
def visit_FunctionDef(self, node):
node.body = self.visit_block(
node.body, after_visit=self._process_body_item)
return node
def test_function(x, y):
z = x
z = y
return z
tr = TestTransformer(self._simple_context())
node, _ = parser.parse_entity(test_function, future_features=())
node = tr.visit(node)
self.assertEqual(len(node.body), 2)
self.assertIsInstance(node.body[0], gast.Assign)
self.assertIsInstance(node.body[1], gast.If)
self.assertIsInstance(node.body[1].body[0], gast.Assign)
self.assertIsInstance(node.body[1].body[1], gast.Return)
def test_robust_error_on_list_visit(self):
class BrokenTransformer(transformer.Base):
def visit_If(self, node):
# This is broken because visit expects a single node, not a list, and
# the body of an if is a list.
# Importantly, the default error handling in visit also expects a single
# node. Therefore, mistakes like this need to trigger a type error
# before the visit called here installs its error handler.
# That type error can then be caught by the enclosing call to visit,
# and correctly blame the If node.
self.visit(node.body)
return node
def test_function(x):
if x > 0:
return x
tr = BrokenTransformer(self._simple_context())
node, _ = parser.parse_entity(test_function, future_features=())
with self.assertRaises(ValueError) as cm:
node = tr.visit(node)
obtained_message = str(cm.exception)
expected_message = r'expected "ast.AST", got "\<(type|class) \'list\'\>"'
self.assertRegex(obtained_message, expected_message)
def test_robust_error_on_ast_corruption(self):
# A child class should not be able to be so broken that it causes the error
# handling in `transformer.Base` to raise an exception. Why not? Because
# then the original error location is dropped, and an error handler higher
# up in the call stack gives misleading information.
# Here we test that the error handling in `visit` completes, and blames the
# correct original exception, even if the AST gets corrupted.
class NotANode(object):
pass
class BrokenTransformer(transformer.Base):
def visit_If(self, node):
node.body = NotANode()
raise ValueError('I blew up')
def test_function(x):
if x > 0:
return x
tr = BrokenTransformer(self._simple_context())
node, _ = parser.parse_entity(test_function, future_features=())
with self.assertRaises(ValueError) as cm:
node = tr.visit(node)
obtained_message = str(cm.exception)
# The message should reference the exception actually raised, not anything
# from the exception handler.
expected_substring = 'I blew up'
self.assertIn(expected_substring, obtained_message)
def test_origin_info_propagated_to_new_nodes(self):
class TestTransformer(transformer.Base):
def visit_If(self, node):
return gast.Pass()
tr = TestTransformer(self._simple_context())
def test_fn():
x = 1
if x > 0:
x = 1
return x
node, source = parser.parse_entity(test_fn, future_features=())
origin_info.resolve(node, source, 'test_file', 100, 0)
node = tr.visit(node)
created_pass_node = node.body[1]
# Takes the line number of the if statement.
self.assertEqual(
anno.getanno(created_pass_node, anno.Basic.ORIGIN).loc.lineno, 102)
def test_origin_info_preserved_in_moved_nodes(self):
class TestTransformer(transformer.Base):
def visit_If(self, node):
return node.body
tr = TestTransformer(self._simple_context())
def test_fn():
x = 1
if x > 0:
x = 1
x += 3
return x
node, source = parser.parse_entity(test_fn, future_features=())
origin_info.resolve(node, source, 'test_file', 100, 0)
node = tr.visit(node)
assign_node = node.body[1]
aug_assign_node = node.body[2]
# Keep their original line numbers.
self.assertEqual(
anno.getanno(assign_node, anno.Basic.ORIGIN).loc.lineno, 103)
self.assertEqual(
anno.getanno(aug_assign_node, anno.Basic.ORIGIN).loc.lineno, 104)
class CodeGeneratorTest(test.TestCase):
def _simple_context(self):
entity_info = transformer.EntityInfo(
name='test_fn',
source_code=None,
source_file=None,
future_features=(),
namespace=None)
return transformer.Context(entity_info, None, None)
def test_basic_codegen(self):
class TestCodegen(transformer.CodeGenerator):
def visit_Assign(self, node):
self.emit(parser.unparse(node, include_encoding_marker=False))
self.emit('\n')
def visit_Return(self, node):
self.emit(parser.unparse(node, include_encoding_marker=False))
self.emit('\n')
def visit_If(self, node):
self.emit('if ')
# This is just for simplifity. A real generator will walk the tree and
# emit proper code.
self.emit(parser.unparse(node.test, include_encoding_marker=False))
self.emit(' {\n')
self.visit_block(node.body)
self.emit('} else {\n')
self.visit_block(node.orelse)
self.emit('}\n')
tg = TestCodegen(self._simple_context())
def test_fn():
x = 1
if x > 0:
x = 2
if x > 1:
x = 3
return x
node, source = parser.parse_entity(test_fn, future_features=())
origin_info.resolve(node, source, 'test_file', 100, 0)
tg.visit(node)
r = re.compile('.*'.join([
r'x = 1',
r'if \(?x > 0\)? {',
r'x = 2',
r'if \(?x > 1\)? {',
r'x = 3',
r'} else {',
r'}',
r'} else {',
r'}',
r'return x']), re.DOTALL)
self.assertRegex(tg.code_buffer, r)
# TODO(mdan): Test the source map.
if __name__ == '__main__':
test.main()
@@ -0,0 +1,496 @@
# Copyright 2016 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.
# ==============================================================================
"""Generic source code transformation infrastructure."""
import inspect
import threading
import types
import gast
from tensorflow.python.autograph.pyct import cache
from tensorflow.python.autograph.pyct import inspect_utils
from tensorflow.python.autograph.pyct import loader
from tensorflow.python.autograph.pyct import naming
from tensorflow.python.autograph.pyct import origin_info
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.autograph.pyct import transformer
from tensorflow.python.autograph.utils import ag_logging as logging
def _wrap_into_factory(nodes, entity_name, inner_factory_name,
outer_factory_name, closure_vars, factory_args,
future_features):
"""Wraps an AST into the body of a factory with consistent lexical context.
The AST is expected to define some symbol with a name given by `entity_name`.
This mechanism ensures that the resulting transformed entity has lexical
scoping identical to that of the source entity, while allowing extra
parametrization.
Two nested factories achieve the following:
1. The inner factory dynamically creates the entity represented by `nodes`.
2. The inner factory is parametrized by a custom set of arguments.
3. The inner factory has a closure identical to that of the transformed
entity.
4. The inner factory has local variables named like `args`, which `nodes` may
use as additional parameters.
5. The inner factory returns the variables given by `entity_name`.
6. The outer factory is niladic.
7. The outer factory has no closure.
8. The outer factory creates the necessary lexical scope for the inner
factory, so that the loaded code has the given configuration for
closure/globals.
9. The outer factory returns the inner factory.
Roughly speaking, the following code is generated:
from __future__ import future_feature_1
from __future__ import future_feature_2
...
def outer_factory():
closure_var_1 = None
closure_var_2 = None
...
def inner_factory(arg_1, arg_2, ...):
<<nodes>>
return entity
return inner_factory
The lexical scoping is created using dummy symbol declarations which create
local variables in the body of the outer factory, so that the Python parser
correctly marks them as free non-global variables upon load (that is, it
creates cell slots for each symbol. These symbols are initialized with None,
but their values are not expected to be used; instead, the caller is expected
to replace them with the cells of the source entity. For more details, see:
https://docs.python.org/3/reference/executionmodel.html#binding-of-names
Args:
nodes: Tuple[ast.AST], the source code to wrap.
entity_name: Union[Text, ast.AST], the name of the principal entity that
`nodes` define.
inner_factory_name: Text, the name of the inner factory.
outer_factory_name: Text, the name of the outer factory.
closure_vars: Iterable[Text], names of the closure variables for the inner
factory.
factory_args: Iterable[Text], names of additional arguments for the
inner factory. Useful to configure variables that the converted code can
use. Typically, these are modules.
future_features: Iterable[Text], names of future statements to associate the
code with.
Returns:
ast.AST
"""
dummy_closure_defs = []
for var_name in closure_vars:
template = """
var_name = None
"""
dummy_closure_defs.extend(templates.replace(template, var_name=var_name))
if future_features:
future_imports = gast.ImportFrom(
module='__future__',
names=[gast.alias(name=name, asname=None) for name in future_features],
level=0)
else:
future_imports = []
factory_args = [
gast.Name(name, ctx=gast.Param(), annotation=None, type_comment=None)
for name in factory_args
]
template = """
future_imports
def outer_factory_name():
dummy_closure_defs
def inner_factory_name(factory_args):
entity_defs
return entity_name
return inner_factory_name
"""
return templates.replace(
template,
dummy_closure_defs=dummy_closure_defs,
entity_defs=nodes,
entity_name=entity_name,
factory_args=factory_args,
future_imports=future_imports,
inner_factory_name=inner_factory_name,
outer_factory_name=outer_factory_name)
class _PythonFnFactory(object):
"""Helper object that wraps a Python function factory."""
def __init__(self, name, freevars, extra_locals):
"""Creates a new factory for a Python function.
Args:
name: The function name.
freevars: The list of non-global free variables for the function.
extra_locals: Dict[Text, Any], names and values for custom variables that
are accessible to the generated code as local variables.
"""
self._name = name
self._freevars = freevars
self._extra_locals = extra_locals
self._unbound_factory = None
self.module = None
self.source_map = None
def create(self,
nodes,
namer,
inner_factory_name='inner_factory',
outer_factory_name='outer_factory',
future_features=()):
"""Initializes a function."""
if self._unbound_factory is not None:
raise ValueError('double initialization; create a new object instead')
inner_factory_name = namer.new_symbol(inner_factory_name, ())
outer_factory_name = namer.new_symbol(outer_factory_name, ())
nodes = _wrap_into_factory(nodes, self._name, inner_factory_name,
outer_factory_name, self._freevars,
self._extra_locals.keys(), future_features)
module, _, source_map = loader.load_ast(
nodes, include_source_map=True)
outer_factory = getattr(module, outer_factory_name)
self._unbound_factory = outer_factory()
self.module = module
self.source_map = source_map
def instantiate(self,
globals_,
closure,
defaults=None,
kwdefaults=None):
"""Creates a new function instance."""
if self._unbound_factory is None:
raise ValueError('call create first')
factory_code = self._unbound_factory.__code__
factory_freevars = factory_code.co_freevars
closure_map = dict(zip(self._freevars, closure))
factory_closure = tuple(
closure_map[name] for name in factory_code.co_freevars)
if len(factory_closure) != len(closure):
raise ValueError(
'closure mismatch, requested {}, but source function had {}'.format(
self._freevars, factory_freevars))
bound_factory = types.FunctionType(
code=factory_code,
globals=globals_,
name=self._name,
argdefs=(),
closure=factory_closure)
# The lint override is a false positive.
new_fn = bound_factory(**self._extra_locals) # pylint:disable=not-callable
if defaults:
new_fn.__defaults__ = defaults
if kwdefaults:
new_fn.__kwdefaults__ = kwdefaults
return new_fn
class GenericTranspiler(object):
"""A generic transpiler for Python functions.
Its interface is the `transform` API, which can process Python function
objects. Internally, it handles parsing.
Users typically subclass this, customizing the `transform_ast` method. The
output of transformed_ast is returned directly by `transform`. Existing
methods like `transform_function` may also be overloaded.
Example:
class MyTransformer(GenericTranspiler):
def transform_ast(self, node, ctx):
result = <<transform node>>
return result
transformer = MyTransformer()
result = transformer.transform(f, ...)
# result is the output
"""
def get_transformed_name(self, node):
"""Returns a name for the output function. Subclasses may override this."""
if isinstance(node, gast.Lambda):
return 'lam'
elif isinstance(node, gast.FunctionDef):
return node.name
raise ValueError('Unknown node type {}'.format(node))
def transform_ast(self, node, ctx):
"""Performs an actual transformation of a function's AST.
Subclasses must implement this method, and do not usually call it.
Args:
node: One or more ast.AST nodes representing the AST to be transformed.
ctx: transformer.Context.
"""
raise NotImplementedError('subclasses must override this')
def transform(self, obj, user_context):
"""Transforms a Python object.
Users typically call this method.
Args:
obj: A Python object, function, type, etc.
user_context: An opaque object (may be None) that is forwarded to
transform_ast, through the ctx.user attribute.
Returns:
The result of calling transform_function.
Raises:
NotImplementedError: if the type of obj is not handled.
"""
if inspect.isfunction(obj) or inspect.ismethod(obj):
return self.transform_function(obj, user_context)
raise NotImplementedError('Non-function: {}'.format(type(obj)))
def _erase_arg_defaults(self, node):
"""Erase arg default expressions, which would otherwise be unbound."""
args = node.args
for i in range(len(args.defaults)):
args.defaults[i] = parser.parse_expression('None')
for i, d in enumerate(args.kw_defaults):
if d is not None:
args.kw_defaults[i] = parser.parse_expression('None')
return node
def transform_module(self, mod, user_context):
"""Transforms a module.
Subclasses may override this method. The return value is opaque.
The method receives the original AST. The result is passed as-is to the
output of `transform`.
Args:
mod: A Python module.
user_context: An opaque object (may be None) that is forwarded to
transform_ast, through the ctx.user attribute.
Returns:
List[Tuple[Any, Any]]. By default it returns the output of transform_ast,
evaluated on each supported member, other than modules, together with a
`transformer.Context` containing information about the transformation
process.
"""
result = []
for member in mod.__dict__.values():
if inspect.ismodule(member):
continue # Not transforming modules recursively.
try:
result.append(self.transform(member, user_context))
except NotImplementedError:
pass # Skip unsupported elements.
return result
def transform_function(self, fn, user_context):
"""Transforms a function.
Subclasses may override this method. The return value is opaque.
The method receives the original AST. The result is passed as-is to the
output of `transform`.
Args:
fn: A function or lambda.
user_context: An opaque object (may be None) that is forwarded to
transform_ast, through the ctx.user attribute.
Returns:
Tuple[Any, Any]. By default it returns the output of transform_ast,
together with a `transformer.Context` containing information about the
transformation process.
"""
future_features = inspect_utils.getfutureimports(fn)
node, source = parser.parse_entity(fn, future_features=future_features)
logging.log(3, 'Source code of %s:\n\n%s\n', fn, source)
origin_info.resolve_entity(node, source, fn)
namespace = inspect_utils.getnamespace(fn)
namer = naming.Namer(namespace)
new_name = namer.new_symbol(self.get_transformed_name(node), ())
entity_info = transformer.EntityInfo(
name=new_name,
source_code=source,
source_file='<fragment>',
future_features=future_features,
namespace=namespace)
context = transformer.Context(entity_info, namer, user_context)
node = self._erase_arg_defaults(node)
result = self.transform_ast(node, context)
return result, context
class PyToPy(GenericTranspiler):
"""A generic Python-to-Python transpiler.
Its `transform` method offers a function-in, function-out interface.
Internally, it takes care of parsing, caching and loading of the translated
code.
Users typically subclass this, overriding `transform_ast`.
Usually, instances of this class are singletons, since each instance manages
its own cache. The caching can be controlled by overriding `get_caching_key`.
Example:
class MyTransformer(PyToPy):
def transform_ast(self, node, ctx):
node = <<transform node, usually using ast.NodeTransformer classes>>
return node
transformer = MyTransformer()
new_f, module, source_map = transformer.transform_function(f, ...)
# new_f is a function with signature identical to f
The transformed function has access to the same namespace as the original
function. To allow access to internal APIs, users may inject additional
symbols by overriding `get_extra_locals`.
"""
def __init__(self):
self._cache_lock = threading.RLock()
self._cache = cache.CodeObjectCache()
def get_extra_locals(self):
"""Returns extra static local variables to be made to transformed code.
Subclasses must override this.
Returns:
extra_locals: A Dict[Text, Any] containing additional variables to make
available to the transformed code.
"""
raise NotImplementedError('subclasses must override this')
def get_caching_key(self, user_context):
"""Returns a unique key to use for caching.
Subclasses must override this.
Calls made to `transform_function` with functions that have the same code
object and caching key will return a cached instance on subsequent
invocations.
Args:
user_context: The context object which was passed to `transform`.
Returns:
extra_locals: A hashable.
"""
raise NotImplementedError('subclasses must override this')
def _cached_factory(self, fn, cache_subkey):
cached_factory = self._cache[fn][cache_subkey]
logging.log(3, 'Cache hit for %s subkey %s: %s', fn, cache_subkey,
cached_factory)
return cached_factory
def transform_function(self, fn, user_context):
"""Transforms a function. See GenericTranspiler.transform_function.
This overload wraps the parent's `transform_function`, adding caching and
facilities to instantiate the output as a Python object. It also
adds facilities to make new symbols available to the generated Python code,
visible as local variables - see `get_extra_locals`.
Args:
fn: A function or lambda.
user_context: An opaque object (may be None) that is forwarded to
transform_ast, through the ctx.user attribute.
Returns:
A tuple:
* A function or lambda with the same signature and closure as `fn`
* The temporary module into which the transformed function was loaded
* The source map as a
Dict[origin_info.LineLocation, origin_info.OriginInfo]
"""
cache_subkey = self.get_caching_key(user_context)
if self._cache.has(fn, cache_subkey):
# Fast path: use a lock-free check.
factory = self._cached_factory(fn, cache_subkey)
else:
with self._cache_lock:
# Check again under lock.
if self._cache.has(fn, cache_subkey):
factory = self._cached_factory(fn, cache_subkey)
else:
logging.log(1, '%s is not cached for subkey %s', fn, cache_subkey)
# TODO(mdan): Confusing overloading pattern. Fix.
nodes, ctx = super(PyToPy, self).transform_function(fn, user_context)
if isinstance(nodes, gast.Lambda):
nodes = gast.Assign(
targets=[
gast.Name(
ctx.info.name,
ctx=gast.Store(),
annotation=None,
type_comment=None)
],
value=nodes)
else:
nodes.name = ctx.info.name
if logging.has_verbosity(2):
logging.log(2, 'Transformed %s:\n\n%s\n', fn, parser.unparse(nodes))
factory = _PythonFnFactory(
ctx.info.name, fn.__code__.co_freevars, self.get_extra_locals())
factory.create(
nodes, ctx.namer, future_features=ctx.info.future_features)
self._cache[fn][cache_subkey] = factory
transformed_fn = factory.instantiate(
globals_=fn.__globals__,
closure=fn.__closure__ or (),
defaults=fn.__defaults__,
kwdefaults=getattr(fn, '__kwdefaults__', None))
return transformed_fn, factory.module, factory.source_map
@@ -0,0 +1,240 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for transpiler module."""
import threading
import gast
from tensorflow.python.autograph.pyct import transformer
from tensorflow.python.autograph.pyct import transpiler
from tensorflow.python.platform import test
class FlipSignTransformer(transformer.Base):
def visit_BinOp(self, node):
if isinstance(node.op, gast.Add):
node.op = gast.Sub()
return self.generic_visit(node)
class TestTranspiler(transpiler.PyToPy):
def get_caching_key(self, ctx):
del ctx
return 0
def get_extra_locals(self):
return {}
def transform_ast(self, node, ctx):
return FlipSignTransformer(ctx).visit(node)
global_var_for_test_global = 1
global_var_for_test_namespace_collisions = object()
class PyToPyTest(test.TestCase):
def test_basic(self):
def f(a):
return a + 1
tr = TestTranspiler()
f, _, _ = tr.transform(f, None)
self.assertEqual(f(1), 0)
def test_closure(self):
b = 1
def f(a):
return a + b
tr = TestTranspiler()
f, _, _ = tr.transform(f, None)
self.assertEqual(f(1), 0)
b = 2
self.assertEqual(f(1), -1)
def test_global(self):
def f(a):
return a + global_var_for_test_global
tr = TestTranspiler()
f, _, _ = tr.transform(f, None)
global global_var_for_test_global
global_var_for_test_global = 1
self.assertEqual(f(1), 0)
global_var_for_test_global = 2
self.assertEqual(f(1), -1)
def test_defaults(self):
b = 2
c = 1
def f(a, d=c + 1):
return a + b + d
tr = TestTranspiler()
f, _, _ = tr.transform(f, None)
self.assertEqual(f(1), 1 - 2 - 2)
c = 0
self.assertEqual(f(1), 1 - 2 - 2) # Defaults are evaluated at definition.
b = 1
self.assertEqual(f(1), 1 - 2 - 1)
def test_call_tree(self):
def g(a):
return a + 1
def f(a):
return g(a) + 1
tr = TestTranspiler()
f, _, _ = tr.transform(f, None)
self.assertEqual(f(1), 1 - 1 + 1) # Only f is converted.
def test_lambda(self):
b = 2
f = lambda x: (b + (x if x > 0 else -x))
tr = TestTranspiler()
f, _, _ = tr.transform(f, None)
self.assertEqual(f(1), 2 - 1)
self.assertEqual(f(-1), 2 - 1)
b = 3
self.assertEqual(f(1), 3 - 1)
self.assertEqual(f(-1), 3 - 1)
def test_multiple_lambdas(self):
a, b = 1, 2
# This can be disambiguated by the argument names.
f, _ = (lambda x: a + x, lambda y: b * y)
tr = TestTranspiler()
f, _, _ = tr.transform(f, None)
self.assertEqual(f(1), 1 - 1)
def test_nested_functions(self):
b = 2
def f(x):
def g(x):
return b + x
return g(x)
tr = TestTranspiler()
f, _, _ = tr.transform(f, None)
self.assertEqual(f(1), 2 - 1)
def test_nested_lambda(self):
b = 2
def f(x):
g = lambda x: b + x
return g(x)
tr = TestTranspiler()
f, _, _ = tr.transform(f, None)
self.assertEqual(f(1), 2 - 1)
def test_concurrency(self):
def f():
pass
outputs = []
tr = TestTranspiler()
# Note: this is not a test, it's a required invariant.
assert tr.get_caching_key(None) == tr.get_caching_key(None)
def conversion_thread():
_, mod, _ = tr.transform(f, None)
outputs.append(mod.__name__)
threads = tuple(
threading.Thread(target=conversion_thread) for _ in range(10))
for t in threads:
t.start()
for t in threads:
t.join()
# Races would potentially create multiple functions / modules
# (non-deterministically, but with high likelihood).
self.assertEqual(len(set(outputs)), 1)
def test_reentrance(self):
def test_fn():
return 1 + 1
class ReentrantTranspiler(transpiler.PyToPy):
def __init__(self):
super(ReentrantTranspiler, self).__init__()
self._recursion_depth = 0
def get_caching_key(self, ctx):
del ctx
return 0
def get_extra_locals(self):
return {}
def transform_ast(self, node, ctx):
self._recursion_depth += 1
if self._recursion_depth < 2:
self.transform(test_fn, None)
return FlipSignTransformer(ctx).visit(node)
tr = ReentrantTranspiler()
f, _, _ = tr.transform(test_fn, None)
self.assertEqual(f(), 0)
def test_namespace_collisions_avoided(self):
class TestClass(object):
def global_var_for_test_namespace_collisions(self):
return global_var_for_test_namespace_collisions
tr = TestTranspiler()
obj = TestClass()
f, _, _ = tr.transform(
obj.global_var_for_test_namespace_collisions, None)
self.assertIs(f(obj), global_var_for_test_namespace_collisions)
if __name__ == '__main__':
test.main()