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()