From d91eef333dc8bf566be901f0c66060ac18bee96d Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Fri, 11 Aug 2023 05:20:23 -0700 Subject: [PATCH] Add optional `node` argument to `get_ast_nodes_of_type`. This enables getting all nodes with an AST type that are also descendants of root node `node`. --- python_graphs/program_graph.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python_graphs/program_graph.py b/python_graphs/program_graph.py index e20093d..bd55cdd 100644 --- a/python_graphs/program_graph.py +++ b/python_graphs/program_graph.py @@ -171,8 +171,12 @@ def get_node_by_ast_node(self, ast_node): def contains_ast_node(self, ast_node): return id(ast_node) in self.ast_id_to_program_graph_node - def get_ast_nodes_of_type(self, ast_type): - for node in six.itervalues(self.nodes): + def get_ast_nodes_of_type(self, ast_type, ast_node=None): + if ast_node: + nodes = self.walk_ast_descendants(ast_node) + else: + nodes = self.nodes + for node in six.itervalues(nodes): if node.node_type == pb.NodeType.AST_NODE and node.ast_type == ast_type: yield node