def test_get_symbols(self):
code = '''
x = 1
d_x[d_idx] += s_x[s_idx]
'''
tree = ast.parse(dedent(code))
result = list(get_symbols(tree))
result.sort()
expect = ['d_idx', 'd_x', 's_idx', 's_x', 'x']
self.assertEqual(result, expect)
# Test if it parses with the code itself instead of a tree.
result = list(get_symbols(dedent(code)))
result.sort()
self.assertEqual(result, expect)
result = list(get_symbols(tree, ctx=ast.Store))
result.sort()
self.assertEqual(result, ['x'])
def test_get_symbols(self):
code = '''
x = 1
d_x[d_idx] += s_x[s_idx]
'''
tree = ast.parse(dedent(code))
result = list(get_symbols(tree))
result.sort()
expect = ['d_idx', 'd_x', 's_idx', 's_x', 'x']
self.assertEqual(result, expect)
# Test if it parses with the code itself instead of a tree.
result = list(get_symbols(dedent(code)))
result.sort()
self.assertEqual(result, expect)
result = list(get_symbols(tree, ctx=ast.Store))
result.sort()
self.assertEqual(result, ['x'])
def _setup_code(self, code):
"""Perform analysis of the code and store the information in various
attributes.
"""
code = dedent(code)
self.code = code
self.ast_tree = ast.parse(code)
self.symbols = get_symbols(self.ast_tree)
symbols = self.symbols
self.src_arrays, self.dest_arrays = get_array_names(symbols)
self._setup_context()
def _setup_code(self, code):
"""Perform analysis of the code and store the information in various
attributes.
"""
code = dedent(code)
self.code = code
self.ast_tree = ast.parse(code)
self.symbols = get_symbols(self.ast_tree)
symbols = self.symbols
self.src_arrays, self.dest_arrays = get_array_names(symbols)
self._setup_context()
