Retrieve all kinds of assignment targets from with test

This commit is contained in:
Batuhan Taskaya
2020-09-24 10:42:56 +03:00
parent 6ae0efa415
commit 606c528803
2 changed files with 11 additions and 2 deletions

View File

@@ -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]

View File

@@ -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"]