xref: /aosp_15_r20/external/tensorflow/tensorflow/python/autograph/pyct/qual_names.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Utilities for manipulating qualified names.
16
17A qualified name is a uniform way to refer to simple (e.g. 'foo') and composite
18(e.g. 'foo.bar') syntactic symbols.
19
20This is *not* related to the __qualname__ attribute used by inspect, which
21refers to scopes.
22"""
23
24import collections
25
26import gast
27
28from tensorflow.python.autograph.pyct import anno
29from tensorflow.python.autograph.pyct import parser
30
31
32class CallerMustSetThis(object):
33  pass
34
35
36class Symbol(collections.namedtuple('Symbol', ['name'])):
37  """Represents a Python symbol."""
38
39
40class Literal(collections.namedtuple('Literal', ['value'])):
41  """Represents a Python numeric literal."""
42
43  def __str__(self):
44    if isinstance(self.value, str):
45      return "'{}'".format(self.value)
46    return str(self.value)
47
48  def __repr__(self):
49    return str(self)
50
51
52# TODO(mdan): Use subclasses to remove the has_attr has_subscript booleans.
53class QN(object):
54  """Represents a qualified name."""
55
56  def __init__(self, base, attr=None, subscript=None):
57    if attr is not None and subscript is not None:
58      raise ValueError('A QN can only be either an attr or a subscript, not '
59                       'both: attr={}, subscript={}.'.format(attr, subscript))
60    self._has_attr = False
61    self._has_subscript = False
62
63    if attr is not None:
64      if not isinstance(base, QN):
65        raise ValueError(
66            'for attribute QNs, base must be a QN; got instead "%s"' % base)
67      if not isinstance(attr, str):
68        raise ValueError('attr may only be a string; got instead "%s"' % attr)
69      self._parent = base
70      # TODO(mdan): Get rid of the tuple - it can only have 1 or 2 elements now.
71      self.qn = (base, attr)
72      self._has_attr = True
73
74    elif subscript is not None:
75      if not isinstance(base, QN):
76        raise ValueError('For subscript QNs, base must be a QN.')
77      self._parent = base
78      self.qn = (base, subscript)
79      self._has_subscript = True
80
81    else:
82      if not isinstance(base, (str, Literal)):
83        # TODO(mdan): Require Symbol instead of string.
84        raise ValueError(
85            'for simple QNs, base must be a string or a Literal object;'
86            ' got instead "%s"' % type(base))
87      assert '.' not in base and '[' not in base and ']' not in base
88      self._parent = None
89      self.qn = (base,)
90
91  def is_symbol(self):
92    return isinstance(self.qn[0], str)
93
94  def is_simple(self):
95    return len(self.qn) <= 1
96
97  def is_composite(self):
98    return len(self.qn) > 1
99
100  def has_subscript(self):
101    return self._has_subscript
102
103  def has_attr(self):
104    return self._has_attr
105
106  @property
107  def attr(self):
108    if not self._has_attr:
109      raise ValueError('Cannot get attr of non-attribute "%s".' % self)
110    return self.qn[1]
111
112  @property
113  def parent(self):
114    if self._parent is None:
115      raise ValueError('Cannot get parent of simple name "%s".' % self.qn[0])
116    return self._parent
117
118  @property
119  def owner_set(self):
120    """Returns all the symbols (simple or composite) that own this QN.
121
122    In other words, if this symbol was modified, the symbols in the owner set
123    may also be affected.
124
125    Examples:
126      'a.b[c.d]' has two owners, 'a' and 'a.b'
127    """
128    owners = set()
129    if self.has_attr() or self.has_subscript():
130      owners.add(self.parent)
131      owners.update(self.parent.owner_set)
132    return owners
133
134  @property
135  def support_set(self):
136    """Returns the set of simple symbols that this QN relies on.
137
138    This would be the smallest set of symbols necessary for the QN to
139    statically resolve (assuming properties and index ranges are verified
140    at runtime).
141
142    Examples:
143      'a.b' has only one support symbol, 'a'
144      'a[i]' has two support symbols, 'a' and 'i'
145    """
146    # TODO(mdan): This might be the set of Name nodes in the AST. Track those?
147    roots = set()
148    if self.has_attr():
149      roots.update(self.parent.support_set)
150    elif self.has_subscript():
151      roots.update(self.parent.support_set)
152      roots.update(self.qn[1].support_set)
153    else:
154      roots.add(self)
155    return roots
156
157  def __hash__(self):
158    return hash(self.qn + (self._has_attr, self._has_subscript))
159
160  def __eq__(self, other):
161    return (isinstance(other, QN) and self.qn == other.qn and
162            self.has_subscript() == other.has_subscript() and
163            self.has_attr() == other.has_attr())
164
165  def __lt__(self, other):
166    return str(self) < str(other)
167
168  def __gt__(self, other):
169    return str(self) > str(other)
170
171  def __str__(self):
172    root = self.qn[0]
173    if self.has_subscript():
174      return '{}[{}]'.format(root, self.qn[1])
175    if self.has_attr():
176      return '.'.join(map(str, self.qn))
177    else:
178      return str(root)
179
180  def __repr__(self):
181    return str(self)
182
183  def ssf(self):
184    """Simple symbol form."""
185    ssfs = [n.ssf() if isinstance(n, QN) else n for n in self.qn]
186    ssf_string = ''
187    for i in range(0, len(self.qn) - 1):
188      if self.has_subscript():
189        delimiter = '_sub_'
190      else:
191        delimiter = '_'
192      ssf_string += ssfs[i] + delimiter
193    return ssf_string + ssfs[-1]
194
195  def ast(self):
196    """AST representation."""
197    # The caller must adjust the context appropriately.
198    if self.has_subscript():
199      return gast.Subscript(
200          value=self.parent.ast(),
201          slice=self.qn[-1].ast(),
202          ctx=CallerMustSetThis)
203    if self.has_attr():
204      return gast.Attribute(
205          value=self.parent.ast(), attr=self.qn[-1], ctx=CallerMustSetThis)
206
207    base = self.qn[0]
208    if isinstance(base, str):
209      return gast.Name(
210          base, ctx=CallerMustSetThis, annotation=None, type_comment=None)
211    elif isinstance(base, Literal):
212      return gast.Constant(base.value, kind=None)
213    else:
214      assert False, ('the constructor should prevent types other than '
215                     'str and Literal')
216
217
218class QnResolver(gast.NodeTransformer):
219  """Annotates nodes with QN information.
220
221  Note: Not using NodeAnnos to avoid circular dependencies.
222  """
223
224  def visit_Name(self, node):
225    node = self.generic_visit(node)
226    anno.setanno(node, anno.Basic.QN, QN(node.id))
227    return node
228
229  def visit_Attribute(self, node):
230    node = self.generic_visit(node)
231    if anno.hasanno(node.value, anno.Basic.QN):
232      anno.setanno(node, anno.Basic.QN,
233                   QN(anno.getanno(node.value, anno.Basic.QN), attr=node.attr))
234    return node
235
236  def visit_Subscript(self, node):
237    # TODO(mdan): This may no longer apply if we overload getitem.
238    node = self.generic_visit(node)
239    s = node.slice
240    if isinstance(s, (gast.Tuple, gast.Slice)):
241      # TODO(mdan): Support range and multi-dimensional indices.
242      # Continuing silently because some demos use these.
243      return node
244    if isinstance(s, gast.Constant) and s.value != Ellipsis:
245      subscript = QN(Literal(s.value))
246    else:
247      # The index may be an expression, case in which a name doesn't make sense.
248      if anno.hasanno(s, anno.Basic.QN):
249        subscript = anno.getanno(s, anno.Basic.QN)
250      else:
251        return node
252    if anno.hasanno(node.value, anno.Basic.QN):
253      anno.setanno(node, anno.Basic.QN,
254                   QN(anno.getanno(node.value, anno.Basic.QN),
255                      subscript=subscript))
256    return node
257
258
259def resolve(node):
260  return QnResolver().visit(node)
261
262
263def from_str(qn_str):
264  node = parser.parse_expression(qn_str)
265  node = resolve(node)
266  return anno.getanno(node, anno.Basic.QN)
267