• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Lint as: python3
2# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15# ==============================================================================
16"""Converting AST to code and Python entities.
17
18Adapted from Tangent.
19"""
20
21from __future__ import absolute_import
22from __future__ import division
23from __future__ import print_function
24
25import atexit
26import errno
27import importlib
28import os
29import sys
30import tempfile
31
32from tensorflow.python.autograph.pyct import origin_info
33from tensorflow.python.autograph.pyct import parser
34
35
36def _remove_file(file_name):
37  """Remove a file, if it exists."""
38  try:
39    os.remove(file_name)
40  except OSError as e:
41    if e.errno == errno.ENOENT:
42      # The file disappeared. Ignore this. Temporary files might get
43      # cleaned up, especially if they reside in /tmp.
44      pass
45    else:
46      raise
47
48
49def load_source(source, delete_on_exit):
50  """Loads the given source code as a Python module."""
51  # TODO(mdan): Drop the linter verride once the CI stops running Py2.
52  with tempfile.NamedTemporaryFile(  # pylint:disable=unexpected-keyword-arg
53      mode='w', suffix='.py', delete=False, encoding='utf-8') as f:
54    module_name = os.path.basename(f.name[:-3])
55    file_name = f.name
56    f.write(source)
57
58  if delete_on_exit:
59    atexit.register(lambda: _remove_file(file_name))
60
61  spec = importlib.util.spec_from_file_location(module_name, file_name)
62  module = importlib.util.module_from_spec(spec)
63  spec.loader.exec_module(module)
64  # TODO(mdan): Use our own garbage-collected cache instead of sys.modules.
65  sys.modules[module_name] = module
66  return module, file_name
67
68
69def load_ast(nodes,
70             indentation='  ',
71             include_source_map=False,
72             delete_on_exit=True):
73  """Loads the given AST as a Python module.
74
75  Compiling the AST code this way ensures that the source code is readable by
76  e.g. `pdb` or `inspect`.
77
78  Args:
79    nodes: Union[ast.AST, Iterable[ast.AST]], the code to compile, as an AST
80      object.
81    indentation: Text, the string to use for indentation.
82    include_source_map: bool, whether return a source map.
83    delete_on_exit: bool, whether to delete the temporary file used for
84      compilation on exit.
85
86  Returns:
87    Tuple[module, Text, Dict[LineLocation, OriginInfo]], containing:
88    the module containing the unparsed nodes, the source code corresponding to
89    nodes, and the source map. Is include_source_map is False, the source map
90    will be None.
91  """
92  if not isinstance(nodes, (list, tuple)):
93    nodes = (nodes,)
94
95  source = parser.unparse(nodes, indentation=indentation)
96  module, _ = load_source(source, delete_on_exit)
97
98  if include_source_map:
99    source_map = origin_info.create_source_map(nodes, source, module.__file__)
100  else:
101    source_map = None
102
103  # TODO(mdan): Return a structured object.
104  return module, source, source_map
105