diff --git a/parso/python/tree.py b/parso/python/tree.py index 55b8058..b6e7fee 100644 --- a/parso/python/tree.py +++ b/parso/python/tree.py @@ -775,8 +775,8 @@ class WithStmt(Flow): return names def get_test_node_from_name(self, name): - node = name.parent - if node.type != 'with_item': + node = search_ancestor(name, "with_item") + if node is None: raise ValueError('The name is not actually part of a with statement.') return node.children[0] diff --git a/test/test_parser_tree.py b/test/test_parser_tree.py index 0201f7c..9a4a2e3 100644 --- a/test/test_parser_tree.py +++ b/test/test_parser_tree.py @@ -229,3 +229,13 @@ def test_iter_funcdefs(): module = parse(code, version='3.8') func_names = [f.name.value for f in module.iter_funcdefs()] assert func_names == ['normal', 'asyn', 'dec_normal', 'dec_async'] + + +def test_with_stmt_get_test_node_from_name(): + code = "with A as X.Y, B as (Z), C as Q[0], D as Q['foo']: pass" + with_stmt = parse(code, version='3').children[0] + tests = [ + with_stmt.get_test_node_from_name(name).value + for name in with_stmt.get_defined_names(include_setitem=True) + ] + assert tests == ["A", "B", "C", "D"]