How to use the asttokens.walk function in asttokens

To help you get started, we’ve selected a few asttokens examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github gristlabs / asttokens / test_asttokens.py View on Github external
def test_mark_tokens_simple(self):
    source = read_fixture('astroid', 'module.py')
    root = ast.parse(source)
    code = asttokens.CodeText(source)
    code.mark_tokens(root)
    all_nodes = list(asttokens.walk(root))

    def get_node_text(line, col, type_name):
      token = code.get_token(line, col)
      for n in all_nodes:
        if n.first_token == token and n.__class__.__name__ == type_name:
          return code.get_text(n)

    # Line 14 is: [indent 4] MY_DICT[key] = val
    self.assertEqual(get_node_text(14, 4, 'Name'), 'MY_DICT')
    self.assertEqual(get_node_text(14, 4, 'Subscript'), 'MY_DICT[key]')
    self.assertEqual(get_node_text(14, 4, 'Assign'), 'MY_DICT[key] = val')

    # Line 35 is: [indent 12] raise XXXError()
    self.assertEqual(get_node_text(35, 12, 'Raise'), 'raise XXXError()')
    self.assertEqual(get_node_text(35, 18, 'Call'), 'XXXError()')
    self.assertEqual(get_node_text(35, 18, 'Name'), 'XXXError')
github gristlabs / asttokens / test_asttokens.py View on Github external
def test_assign_first_tokens(self):
    source = read_fixture('astroid', 'module.py')
    root = ast.parse(source)
    code = asttokens.CodeText(source)
    code.mark_tokens(root)
    all_nodes = list(asttokens.walk(root))

    def get_node_types_at(line, col):
      token = code.get_token(line, col)
      return {n.__class__.__name__ for n in all_nodes if n.first_token == token}

    # Line 14 is: [indent 4] MY_DICT[key] = val
    self.assertEqual(get_node_types_at(14, 4), {'Name', 'Subscript', 'Assign'})

    # Line 35 is: [indent 12] raise XXXError()
    self.assertEqual(get_node_types_at(35, 12), {'Raise'})
    self.assertEqual(get_node_types_at(35, 18), {'Call', 'Name'})

    # Line 53 is: [indent 12] autre = [a for (a, b) in MY_DICT if b]
    self.assertEqual(get_node_types_at(53, 20), {'ListComp'})
    self.assertEqual(get_node_types_at(53, 21), {'Name'})
    self.assertEqual(get_node_types_at(53, 23), {'comprehension'})