Rename variable

Author: stellasia

src/neo4j_graphrag/experimental/components/graph_pruning.py

@@ -223,21 +223,21 @@ def _validate_node(
 
     def _enforce_nodes(
         self,
-        extracted_nodes: list[Neo4jNode],
+        nodes: list[Neo4jNode],
         schema: GraphSchema,
         lexical_graph_config: LexicalGraphConfig,
         pruning_stats: PruningStats,
     ) -> list[Neo4jNode]:
         """
-        Filter extracted nodes to be conformant to the schema.
+        Filter nodes to be conformant to the schema.
 
         Keep only those whose label is in schema
         (unless schema has additional_node_types=True, default value)
         For each valid node, validate properties. If a node is left without
         properties, prune it.
         """
         valid_nodes = []
-        for node in extracted_nodes:
+        for node in nodes:
             if node.label in lexical_graph_config.lexical_graph_node_labels:
                 valid_nodes.append(node)
                 continue
@@ -329,14 +329,14 @@ def _validate_relationship(
 
     def _enforce_relationships(
         self,
-        extracted_relationships: list[Neo4jRelationship],
+        relationships: list[Neo4jRelationship],
         filtered_nodes: list[Neo4jNode],
         schema: GraphSchema,
         lexical_graph_config: LexicalGraphConfig,
         pruning_stats: PruningStats,
     ) -> list[Neo4jRelationship]:
         """
-        Filter extracted nodes to be conformant to the schema.
+        Filter relationships to be conformant to the schema.
 
         Keep only those whose types are in schema, start/end node conform to schema,
         and start/end nodes are in filtered nodes (i.e., kept after node enforcement).
@@ -347,7 +347,7 @@ def _enforce_relationships(
 
         valid_rels = []
         valid_nodes = {node.id: node.label for node in filtered_nodes}
-        for rel in extracted_relationships:
+        for rel in relationships:
             if rel.type in lexical_graph_config.lexical_graph_relationship_types:
                 valid_rels.append(rel)
                 continue

tests/unit/experimental/components/test_graph_pruning.py

@@ -172,7 +172,7 @@ def test_graph_pruning_enforce_nodes_lexical_graph(
 ) -> None:
     pruner = GraphPruning()
     result = pruner._enforce_nodes(
-        extracted_nodes=[
+        nodes=[
             Neo4jNode(id="1", label="Paragraph"),
         ],
         schema=GraphSchema(node_types=tuple(), additional_node_types=False),