make gc ast bookkeeping optional

Author: smacke

pyccolo/ast_rewriter.py

@@ -39,6 +39,8 @@
 
 
 class AstRewriter(ast.NodeTransformer):
+    gc_bookkeeping = True
+
     def __init__(
         self,
         tracers: "List[BaseTracer]",
@@ -112,7 +114,7 @@ def visit(self, node: ast.AST):
         new_bookkeeper = last_tracer.ast_bookkeeper_by_fname[
             self._path
         ] = AstBookkeeper.create(self._path, module_id)
-        if old_bookkeeper is not None:
+        if old_bookkeeper is not None and self.gc_bookkeeping:
             last_tracer.remove_bookkeeping(old_bookkeeper, module_id)
         BookkeepingVisitor(
             new_bookkeeper.ast_node_by_id,

pyccolo/tracer.py

@@ -207,14 +207,17 @@ def __init__(self, is_reset: bool = False):
         self._post_init_hook_start()
 
     @classmethod
-    def remove_bookkeeping(cls, bookkeeper: AstBookkeeper, module_id: int) -> None:
+    def remove_bookkeeping(
+        cls, bookkeeper: AstBookkeeper, module_id: Optional[int]
+    ) -> None:
         clear_keys(cls.ast_node_by_id, bookkeeper.ast_node_by_id)
         clear_keys(cls.containing_ast_by_id, bookkeeper.containing_ast_by_id)
         clear_keys(cls.containing_stmt_by_id, bookkeeper.containing_stmt_by_id)
         clear_keys(cls.parent_stmt_by_id, bookkeeper.parent_stmt_by_id)
-        clear_keys(
-            cls.stmt_by_lineno_by_module_id[module_id], bookkeeper.stmt_by_lineno
-        )
+        if module_id is not None:
+            clear_keys(
+                cls.stmt_by_lineno_by_module_id[module_id], bookkeeper.stmt_by_lineno
+            )
         for spec, node_ids in cls.augmented_node_ids_by_spec.items():
             clear_keys(node_ids, bookkeeper.ast_node_by_id)