From 606c528803f43ec9fd607c4e0c53e4d2db004eb9 Mon Sep 17 00:00:00 2001 From: Batuhan Taskaya Date: Thu, 24 Sep 2020 10:42:56 +0300 Subject: [PATCH] Retrieve all kinds of assignment targets from with test --- parso/python/tree.py | 4 ++-- test/test_parser_tree.py | 9 +++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) 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..dfe41ca 100644 --- a/test/test_parser_tree.py +++ b/test/test_parser_tree.py @@ -229,3 +229,12 @@ 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"]