From 7c88edf5450fa808d59e48f0a19603127e2a8a58 Mon Sep 17 00:00:00 2001
From: Daniel Flook <daniel@flook.org>
Date: Fri, 18 Oct 2024 13:47:48 +0100
Subject: [PATCH] Use free function to get a nodes parent.

Creates a new module for getting and setting the parent node
---
 .github/workflows/release_test.yaml           |  2 +-
 src/python_minifier/__init__.py               |  2 +
 .../ast_annotation/__init__.py                | 82 +++++++++++++++++++
 src/python_minifier/rename/mapper.py          | 73 ++++++++---------
 src/python_minifier/rename/rename_literals.py |  7 +-
 .../transforms/constant_folding.py            |  3 +-
 .../transforms/remove_annotations.py          | 19 +++--
 .../transforms/remove_exception_brackets.py   | 11 +--
 .../transforms/suite_transformer.py           |  6 +-
 test/ast_annotation/test_add_parent.py        | 49 +++++++++++
 test/helpers.py                               |  2 +
 test/test_combine_imports.py                  |  3 +
 test/test_folding.py                          |  2 +
 test/test_hoist_literals.py                   |  2 +
 test/test_match_rename.py                     |  2 +
 test/test_remove_annotations.py               |  2 +
 test/test_remove_assert.py                    |  2 +
 test/test_remove_debug.py                     |  2 +
 test/test_remove_exception_brackets.py        |  2 +
 test/test_remove_literal_statements.py        |  2 +
 test/test_remove_pass.py                      |  2 +
 test/test_rename_builtins.py                  |  2 +
 test/test_rename_locals.py                    |  2 +
 23 files changed, 221 insertions(+), 60 deletions(-)
 create mode 100644 src/python_minifier/ast_annotation/__init__.py
 create mode 100644 test/ast_annotation/test_add_parent.py

diff --git a/.github/workflows/release_test.yaml b/.github/workflows/release_test.yaml
index 94d7534b..241f4f32 100644
--- a/.github/workflows/release_test.yaml
+++ b/.github/workflows/release_test.yaml
@@ -188,7 +188,7 @@ jobs:
 
       - name: Test typing
         run: |
-          pip3.13 install mypy!=1.12.0 types-setuptools
+          pip3.13 install 'mypy<1.12.0' types-setuptools
           mypy --strict typing_test/test_typing.py
 
           if mypy --strict typing_test/test_badtyping.py; then
diff --git a/src/python_minifier/__init__.py b/src/python_minifier/__init__.py
index 7c6b35f3..15744f1f 100644
--- a/src/python_minifier/__init__.py
+++ b/src/python_minifier/__init__.py
@@ -7,6 +7,7 @@
 import re
 
 import python_minifier.ast_compat as ast
+from python_minifier.ast_annotation import add_parent
 
 from python_minifier.ast_compare import CompareError, compare_ast
 from python_minifier.module_printer import ModulePrinter
@@ -115,6 +116,7 @@ def minify(
     # This will raise if the source file can't be parsed
     module = ast.parse(source, filename)
 
+    add_parent(module)
     add_namespace(module)
 
     if remove_literal_statements:
diff --git a/src/python_minifier/ast_annotation/__init__.py b/src/python_minifier/ast_annotation/__init__.py
new file mode 100644
index 00000000..9cf7a35c
--- /dev/null
+++ b/src/python_minifier/ast_annotation/__init__.py
@@ -0,0 +1,82 @@
+"""
+This module provides utilities for annotating Abstract Syntax Tree (AST) nodes with parent references.
+"""
+
+import ast
+
+class _NoParent(ast.AST):
+    """A placeholder class used to indicate that a node has no parent."""
+
+    def __repr__(self):
+        # type: () -> str
+        return 'NoParent()'
+
+def add_parent(node, parent=_NoParent()):
+    # type: (ast.AST, ast.AST) -> None
+    """
+    Recursively adds a parent reference to each node in the AST.
+
+    >>> tree = ast.parse('a = 1')
+    >>> add_parent(tree)
+    >>> get_parent(tree.body[0]) == tree
+    True
+
+    :param node: The current AST node.
+    :param parent: The parent :class:`ast.AST` node.
+    """
+
+    node._parent = parent  # type: ignore[attr-defined]
+    for child in ast.iter_child_nodes(node):
+        add_parent(child, node)
+
+def get_parent(node):
+    # type: (ast.AST) -> ast.AST
+    """
+    Retrieves the parent of the given AST node.
+
+    >>> tree = ast.parse('a = 1')
+    >>> add_parent(tree)
+    >>> get_parent(tree.body[0]) == tree
+    True
+
+    :param node: The AST node whose parent is to be retrieved.
+    :return: The parent AST node.
+    :raises ValueError: If the node has no parent.
+    """
+
+    if not hasattr(node, '_parent') or isinstance(node._parent, _NoParent):  # type: ignore[attr-defined]
+        raise ValueError('Node has no parent')
+
+    return node._parent  # type: ignore[attr-defined]
+
+def set_parent(node, parent):
+    # type: (ast.AST, ast.AST) -> None
+    """
+    Replace the parent of the given AST node.
+
+    Create a simple AST:
+    >>> tree = ast.parse('a = func()')
+    >>> add_parent(tree)
+    >>> isinstance(tree.body[0], ast.Assign) and isinstance(tree.body[0].value, ast.Call)
+    True
+    >>> assign = tree.body[0]
+    >>> call = tree.body[0].value
+    >>> get_parent(call) == assign
+    True
+
+    Replace the parent of the call node:
+    >>> tree.body[0] = call
+    >>> set_parent(call, tree)
+    >>> get_parent(call) == tree
+    True
+    >>> from python_minifier.ast_printer import print_ast
+    >>> print(print_ast(tree))
+    Module(body=[
+        Call(Name('func'))
+    ])
+
+    :param node: The AST node whose parent is to be set.
+    :param parent: The parent AST node.
+    """
+
+    node._parent = parent  # type: ignore[attr-defined]
diff --git a/src/python_minifier/rename/mapper.py b/src/python_minifier/rename/mapper.py
index 5b2d8150..44567867 100644
--- a/src/python_minifier/rename/mapper.py
+++ b/src/python_minifier/rename/mapper.py
@@ -3,47 +3,47 @@
 """
 
 import python_minifier.ast_compat as ast
+from python_minifier.ast_annotation import get_parent
 
 from python_minifier.rename.util import is_namespace
 
 
 def add_parent_to_arguments(arguments, func):
-    arguments.parent = func
     arguments.namespace = func
 
     for arg in getattr(arguments, 'posonlyargs', []) + arguments.args:
-        add_parent(arg, arguments, func)
+        add_parent(arg, func)
         if hasattr(arg, 'annotation') and arg.annotation is not None:
-            add_parent(arg.annotation, arguments, func.namespace)
+            add_parent(arg.annotation, func.namespace)
 
     if hasattr(arguments, 'kwonlyargs'):
         for arg in arguments.kwonlyargs:
-            add_parent(arg, arguments, func)
+            add_parent(arg, func)
             if arg.annotation is not None:
-                add_parent(arg.annotation, arguments, func.namespace)
+                add_parent(arg.annotation, func.namespace)
 
         for node in arguments.kw_defaults:
             if node is not None:
-                add_parent(node, arguments, func.namespace)
+                add_parent(node, func.namespace)
 
     for node in arguments.defaults:
-        add_parent(node, arguments, func.namespace)
+        add_parent(node, func.namespace)
 
     if arguments.vararg:
         if hasattr(arguments, 'varargannotation') and arguments.varargannotation is not None:
-            add_parent(arguments.varargannotation, arguments, func.namespace)
+            add_parent(arguments.varargannotation, func.namespace)
         elif isinstance(arguments.vararg, str):
             pass
         else:
-            add_parent(arguments.vararg, arguments, func)
+            add_parent(arguments.vararg, func)
 
     if arguments.kwarg:
         if hasattr(arguments, 'kwargannotation') and arguments.kwargannotation is not None:
-            add_parent(arguments.kwargannotation, arguments, func.namespace)
+            add_parent(arguments.kwargannotation, func.namespace)
         elif isinstance(arguments.kwarg, str):
             pass
         else:
-            add_parent(arguments.kwarg, arguments, func)
+            add_parent(arguments.kwarg, func)
 
 
 def add_parent_to_functiondef(functiondef):
@@ -55,17 +55,17 @@ def add_parent_to_functiondef(functiondef):
         add_parent_to_arguments(functiondef.args, func=functiondef)
 
     for node in functiondef.body:
-        add_parent(node, parent=functiondef, namespace=functiondef)
+        add_parent(node, namespace=functiondef)
 
     for node in functiondef.decorator_list:
-        add_parent(node, parent=functiondef, namespace=functiondef.namespace)
+        add_parent(node, namespace=functiondef.namespace)
 
     if hasattr(functiondef, 'type_params') and functiondef.type_params is not None:
         for node in functiondef.type_params:
-            add_parent(node, parent=functiondef, namespace=functiondef.namespace)
+            add_parent(node, namespace=functiondef.namespace)
 
     if hasattr(functiondef, 'returns') and functiondef.returns is not None:
-        add_parent(functiondef.returns, parent=functiondef, namespace=functiondef.namespace)
+        add_parent(functiondef.returns, namespace=functiondef.namespace)
 
 
 def add_parent_to_classdef(classdef):
@@ -74,65 +74,60 @@ def add_parent_to_classdef(classdef):
     """
 
     for node in classdef.bases:
-        add_parent(node, parent=classdef, namespace=classdef.namespace)
+        add_parent(node, namespace=classdef.namespace)
 
     if hasattr(classdef, 'keywords'):
         for node in classdef.keywords:
-            add_parent(node, parent=classdef, namespace=classdef.namespace)
+            add_parent(node, namespace=classdef.namespace)
 
     if hasattr(classdef, 'starargs') and classdef.starargs is not None:
-        add_parent(classdef.starargs, parent=classdef, namespace=classdef.namespace)
+        add_parent(classdef.starargs, namespace=classdef.namespace)
 
     if hasattr(classdef, 'kwargs') and classdef.kwargs is not None:
-        add_parent(classdef.kwargs, parent=classdef, namespace=classdef.namespace)
+        add_parent(classdef.kwargs, namespace=classdef.namespace)
 
     for node in classdef.body:
-        add_parent(node, parent=classdef, namespace=classdef)
+        add_parent(node, namespace=classdef)
 
     for node in classdef.decorator_list:
-        add_parent(node, parent=classdef, namespace=classdef.namespace)
+        add_parent(node, namespace=classdef.namespace)
 
     if hasattr(classdef, 'type_params') and classdef.type_params is not None:
         for node in classdef.type_params:
-            add_parent(node, parent=classdef, namespace=classdef.namespace)
+            add_parent(node, namespace=classdef.namespace)
 
 
 def add_parent_to_comprehension(node, namespace):
     assert isinstance(node, (ast.GeneratorExp, ast.SetComp, ast.DictComp, ast.ListComp))
 
     if hasattr(node, 'elt'):
-        add_parent(node.elt, parent=node, namespace=node)
+        add_parent(node.elt, namespace=node)
     elif hasattr(node, 'key'):
-        add_parent(node.key, parent=node, namespace=node)
-        add_parent(node.value, parent=node, namespace=node)
+        add_parent(node.key, namespace=node)
+        add_parent(node.value, namespace=node)
 
     iter_namespace = namespace
     for generator in node.generators:
-        generator.parent = node
         generator.namespace = node
 
-        add_parent(generator.target, parent=generator, namespace=node)
-        add_parent(generator.iter, parent=generator, namespace=iter_namespace)
+        add_parent(generator.target, namespace=node)
+        add_parent(generator.iter, namespace=iter_namespace)
         iter_namespace = node
         for if_ in generator.ifs:
-            add_parent(if_, parent=generator, namespace=node)
+            add_parent(if_, namespace=node)
 
 
-def add_parent(node, parent=None, namespace=None):
+def add_parent(node, namespace=None):
     """
-    Add a parent attribute to child nodes
     Add a namespace attribute to child nodes
 
-    :param node: The tree to add parent and namespace properties to
+    :param node: The tree to add namespace properties to
     :type node: :class:`ast.AST`
-    :param parent: The parent node of this node
-    :type parent: :class:`ast.AST`
     :param namespace: The namespace Node that this node is in
     :type namespace: ast.Lambda or ast.Module or ast.FunctionDef or ast.AsyncFunctionDef or ast.ClassDef or ast.DictComp or ast.SetComp or ast.ListComp or ast.Generator
 
     """
 
-    node.parent = parent if parent is not None else node
     node.namespace = namespace if namespace is not None else node
 
     if is_namespace(node):
@@ -146,12 +141,12 @@ def add_parent(node, parent=None, namespace=None):
             add_parent_to_comprehension(node, namespace=namespace)
         elif isinstance(node, ast.Lambda):
             add_parent_to_arguments(node.args, func=node)
-            add_parent(node.body, parent=node, namespace=node)
+            add_parent(node.body, namespace=node)
         elif isinstance(node, ast.ClassDef):
             add_parent_to_classdef(node)
         else:
             for child in ast.iter_child_nodes(node):
-                add_parent(child, parent=node, namespace=node)
+                add_parent(child, namespace=node)
 
         return
 
@@ -163,11 +158,11 @@ def add_parent(node, parent=None, namespace=None):
     if isinstance(node, ast.Name) and isinstance(namespace, ast.ClassDef):
         if isinstance(node.ctx, ast.Load):
             namespace.nonlocal_names.add(node.id)
-        elif isinstance(node.ctx, ast.Store) and isinstance(node.parent, ast.AugAssign):
+        elif isinstance(node.ctx, ast.Store) and isinstance(get_parent(node), ast.AugAssign):
             namespace.nonlocal_names.add(node.id)
 
     for child in ast.iter_child_nodes(node):
-        add_parent(child, parent=node, namespace=namespace)
+        add_parent(child, namespace=namespace)
 
 
 def add_namespace(module):
diff --git a/src/python_minifier/rename/rename_literals.py b/src/python_minifier/rename/rename_literals.py
index 59e1caea..c69abd95 100644
--- a/src/python_minifier/rename/rename_literals.py
+++ b/src/python_minifier/rename/rename_literals.py
@@ -1,4 +1,5 @@
 import python_minifier.ast_compat as ast
+from python_minifier.ast_annotation import get_parent, set_parent
 
 from python_minifier.rename.binding import Binding
 from python_minifier.rename.util import insert
@@ -7,8 +8,8 @@
 
 
 def replace(old_node, new_node):
-    parent = old_node.parent
-    new_node.parent = parent
+    parent = get_parent(old_node)
+    set_parent(new_node, parent)
     new_node.namespace = old_node.namespace
 
     for field, old_value in ast.iter_fields(parent):
@@ -202,7 +203,7 @@ def get_binding(self, value, node):
 
     def visit_Str(self, node):
 
-        if isinstance(node.parent, ast.Expr):
+        if isinstance(get_parent(node), ast.Expr):
             # This is literal statement
             # The RemoveLiteralStatements transformer must have left it here, so ignore it.
             return
diff --git a/src/python_minifier/transforms/constant_folding.py b/src/python_minifier/transforms/constant_folding.py
index e21dcc86..ec041eb2 100644
--- a/src/python_minifier/transforms/constant_folding.py
+++ b/src/python_minifier/transforms/constant_folding.py
@@ -2,6 +2,7 @@
 import sys
 
 import python_minifier.ast_compat as ast
+from python_minifier.ast_annotation import get_parent
 
 from python_minifier.ast_compare import compare_ast
 from python_minifier.expression_printer import ExpressionPrinter
@@ -93,7 +94,7 @@ def visit_BinOp(self, node):
             return node
 
         # New representation is shorter and has the same value, so use it
-        return self.add_child(new_node, node.parent, node.namespace)
+        return self.add_child(new_node, get_parent(node), node.namespace)
 
 
 def equal_value_and_type(a, b):
diff --git a/src/python_minifier/transforms/remove_annotations.py b/src/python_minifier/transforms/remove_annotations.py
index 274d7ea8..8ac6e6b9 100644
--- a/src/python_minifier/transforms/remove_annotations.py
+++ b/src/python_minifier/transforms/remove_annotations.py
@@ -1,6 +1,7 @@
 import sys
 
 import python_minifier.ast_compat as ast
+from python_minifier.ast_annotation import get_parent
 
 from python_minifier.transforms.remove_annotations_options import RemoveAnnotationsOptions
 from python_minifier.transforms.suite_transformer import SuiteTransformer
@@ -72,13 +73,13 @@ def is_dataclass_field(node):
             if sys.version_info < (3, 7):
                 return False
 
-            if not isinstance(node.parent, ast.ClassDef):
+            if not isinstance(get_parent(node), ast.ClassDef):
                 return False
 
-            if len(node.parent.decorator_list) == 0:
+            if len(get_parent(node).decorator_list) == 0:
                 return False
 
-            for decorator_node in node.parent.decorator_list:
+            for decorator_node in get_parent(node).decorator_list:
                 if isinstance(decorator_node, ast.Name) and decorator_node.id == 'dataclass':
                     return True
                 elif isinstance(decorator_node, ast.Attribute) and decorator_node.attr == 'dataclass':
@@ -94,15 +95,15 @@ def is_typing_sensitive(node):
             if sys.version_info < (3, 5):
                 return False
 
-            if not isinstance(node.parent, ast.ClassDef):
+            if not isinstance(get_parent(node), ast.ClassDef):
                 return False
 
-            if len(node.parent.bases) == 0:
+            if len(get_parent(node).bases) == 0:
                 return False
 
             tricky_types = ['NamedTuple', 'TypedDict']
 
-            for base_node in node.parent.bases:
+            for base_node in get_parent(node).bases:
                 if isinstance(base_node, ast.Name) and base_node.id in tricky_types:
                     return True
                 elif isinstance(base_node, ast.Attribute) and base_node.attr in tricky_types:
@@ -111,7 +112,7 @@ def is_typing_sensitive(node):
             return False
 
         # is this a class attribute or a variable?
-        if isinstance(node.parent, ast.ClassDef):
+        if isinstance(get_parent(node), ast.ClassDef):
             if not self._options.remove_class_attribute_annotations:
                 return node
         else:
@@ -121,11 +122,11 @@ def is_typing_sensitive(node):
         if is_dataclass_field(node) or is_typing_sensitive(node):
             return node
         elif node.value:
-            return self.add_child(ast.Assign([node.target], node.value), parent=node.parent, namespace=node.namespace)
+            return self.add_child(ast.Assign([node.target], node.value), parent=get_parent(node), namespace=node.namespace)
         else:
             # Valueless annotations cause the interpreter to treat the variable as a local.
             # I don't know of another way to do that without assigning to it, so
             # keep it as an AnnAssign, but replace the annotation with '0'
 
-            node.annotation = self.add_child(ast.Num(0), parent=node.parent, namespace=node.namespace)
+            node.annotation = self.add_child(ast.Num(0), parent=get_parent(node), namespace=node.namespace)
             return node
diff --git a/src/python_minifier/transforms/remove_exception_brackets.py b/src/python_minifier/transforms/remove_exception_brackets.py
index fedc59b1..0b5df902 100644
--- a/src/python_minifier/transforms/remove_exception_brackets.py
+++ b/src/python_minifier/transforms/remove_exception_brackets.py
@@ -11,6 +11,7 @@
 import sys
 
 import python_minifier.ast_compat as ast
+from python_minifier.ast_annotation import get_parent, set_parent
 
 from python_minifier.rename.binding import BuiltinBinding
 
@@ -84,15 +85,15 @@ def _remove_empty_call(binding):
         assert isinstance(name_node, ast.Name)
         assert isinstance(name_node.ctx, ast.Load)
 
-        if not isinstance(name_node.parent, ast.Call):
+        if not isinstance(get_parent(name_node), ast.Call):
             # This is not a call
             continue
-        call_node = name_node.parent
+        call_node = get_parent(name_node)
 
-        if not isinstance(call_node.parent, ast.Raise):
+        if not isinstance(get_parent(call_node), ast.Raise):
             # This is not a raise statement
             continue
-        raise_node = call_node.parent
+        raise_node = get_parent(call_node)
 
         if len(call_node.args) > 0 or len(call_node.keywords) > 0:
             # This is a call with arguments
@@ -105,7 +106,7 @@ def _remove_empty_call(binding):
             raise_node.exc = name_node
         elif raise_node.cause is call_node:
             raise_node.cause = name_node
-        name_node.parent = raise_node
+        set_parent(name_node, raise_node)
 
 
 def remove_no_arg_exception_call(module):
diff --git a/src/python_minifier/transforms/suite_transformer.py b/src/python_minifier/transforms/suite_transformer.py
index 9226c2fc..b66afa2d 100644
--- a/src/python_minifier/transforms/suite_transformer.py
+++ b/src/python_minifier/transforms/suite_transformer.py
@@ -1,4 +1,5 @@
 import python_minifier.ast_compat as ast
+from python_minifier.ast_annotation import get_parent, add_parent as add_node_parent
 
 from python_minifier.rename.mapper import add_parent
 
@@ -186,10 +187,11 @@ def nearest_function_namespace(node):
 
             if isinstance(node, (ast.FunctionDef, ast.Module, ast.AsyncFunctionDef)):
                 return node
-            return nearest_function_namespace(node.parent)
+            return nearest_function_namespace(get_parent(node))
 
         if namespace is None:
             namespace = nearest_function_namespace(parent)
 
-        add_parent(child, parent=parent, namespace=namespace)
+        add_node_parent(child, parent=parent)
+        add_parent(child, namespace=namespace)
         return child
diff --git a/test/ast_annotation/test_add_parent.py b/test/ast_annotation/test_add_parent.py
new file mode 100644
index 00000000..83e434ee
--- /dev/null
+++ b/test/ast_annotation/test_add_parent.py
@@ -0,0 +1,49 @@
+import pytest
+import ast
+from python_minifier.ast_annotation import add_parent, get_parent, set_parent
+
+
+def test_add_parent():
+
+    source = '''
+class A:
+    def b(self):
+        pass
+'''
+
+    tree = ast.parse(source)
+
+    add_parent(tree)
+
+    assert isinstance(tree, ast.Module)
+
+    assert isinstance(tree.body[0], ast.ClassDef)
+    assert get_parent(tree.body[0]) is tree
+
+    assert isinstance(tree.body[0].body[0], ast.FunctionDef)
+    assert get_parent(tree.body[0].body[0]) is tree.body[0]
+
+    assert isinstance(tree.body[0].body[0].body[0], ast.Pass)
+    assert get_parent(tree.body[0].body[0].body[0]) is tree.body[0].body[0]
+
+
+def test_no_parent_for_root_node():
+    tree = ast.parse('a = 1')
+    add_parent(tree)
+    with pytest.raises(ValueError):
+        get_parent(tree)
+
+
+def test_no_parent_for_unannotated_node():
+    tree = ast.parse('a = 1')
+    with pytest.raises(ValueError):
+        get_parent(tree.body[0])
+
+
+def test_replaces_parent_of_given_node():
+    tree = ast.parse('a = func()')
+    add_parent(tree)
+    call = tree.body[0].value
+    tree.body[0] = call
+    set_parent(call, tree)
+    assert get_parent(call) == tree
diff --git a/test/helpers.py b/test/helpers.py
index c88c7a09..b1452138 100644
--- a/test/helpers.py
+++ b/test/helpers.py
@@ -1,4 +1,5 @@
 import python_minifier.ast_compat as ast
+from python_minifier.ast_annotation import add_parent
 
 from python_minifier.rename import add_namespace, resolve_names
 from python_minifier.rename.bind_names import bind_names
@@ -9,6 +10,7 @@
 def assert_namespace_tree(source, expected_tree):
     tree = ast.parse(source)
 
+    add_parent(tree)
     add_namespace(tree)
     bind_names(tree)
     resolve_names(tree)
diff --git a/test/test_combine_imports.py b/test/test_combine_imports.py
index 2b8081d8..9f08f604 100644
--- a/test/test_combine_imports.py
+++ b/test/test_combine_imports.py
@@ -1,6 +1,7 @@
 import ast
 
 from helpers import print_namespace
+from python_minifier.ast_annotation import add_parent
 
 from python_minifier.ast_compare import compare_ast
 from python_minifier.rename import add_namespace, bind_names, resolve_names
@@ -8,6 +9,7 @@
 
 
 def combine_imports(module):
+    add_parent(module)
     add_namespace(module)
     CombineImports()(module)
     return module
@@ -16,6 +18,7 @@ def combine_imports(module):
 def assert_namespace_tree(source, expected_tree):
     tree = ast.parse(source)
 
+    add_parent(tree)
     add_namespace(tree)
     CombineImports()(tree)
     bind_names(tree)
diff --git a/test/test_folding.py b/test/test_folding.py
index f4004a90..7a8af44c 100644
--- a/test/test_folding.py
+++ b/test/test_folding.py
@@ -3,6 +3,7 @@
 
 import pytest
 
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import compare_ast
 from python_minifier.rename import add_namespace
 from python_minifier.transforms.constant_folding import FoldConstants
@@ -10,6 +11,7 @@
 
 def fold_constants(source):
     module = ast.parse(source)
+    add_parent(module)
     add_namespace(module)
     FoldConstants()(module)
     return module
diff --git a/test/test_hoist_literals.py b/test/test_hoist_literals.py
index 870b7207..3a75e2ac 100644
--- a/test/test_hoist_literals.py
+++ b/test/test_hoist_literals.py
@@ -4,6 +4,7 @@
 import pytest
 
 from python_minifier import unparse
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import compare_ast
 from python_minifier.ast_printer import print_ast
 from python_minifier.rename import (
@@ -19,6 +20,7 @@
 
 def hoist(source):
     module = ast.parse(source)
+    add_parent(module)
     add_namespace(module)
     bind_names(module)
     resolve_names(module)
diff --git a/test/test_match_rename.py b/test/test_match_rename.py
index ef1a088a..89860a99 100644
--- a/test/test_match_rename.py
+++ b/test/test_match_rename.py
@@ -4,6 +4,7 @@
 import pytest
 
 from python_minifier import unparse
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import CompareError, compare_ast
 from python_minifier.rename import (
     add_namespace,
@@ -19,6 +20,7 @@
 def do_rename(source, allow_locals, allow_globals):
     # This will raise if the source file can't be parsed
     module = ast.parse(source, 'test_match_rename')
+    add_parent(module)
     add_namespace(module)
     bind_names(module)
     resolve_names(module)
diff --git a/test/test_remove_annotations.py b/test/test_remove_annotations.py
index e61af0ae..a0b4e61c 100644
--- a/test/test_remove_annotations.py
+++ b/test/test_remove_annotations.py
@@ -4,6 +4,7 @@
 import pytest
 
 from python_minifier import RemoveAnnotationsOptions
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import compare_ast
 from python_minifier.rename import add_namespace
 from python_minifier.transforms.remove_annotations import RemoveAnnotations
@@ -11,6 +12,7 @@
 
 def remove_annotations(source, **kwargs):
     module = ast.parse(source)
+    add_parent(module)
     add_namespace(module)
     RemoveAnnotations(RemoveAnnotationsOptions(**kwargs))(module)
     return module
diff --git a/test/test_remove_assert.py b/test/test_remove_assert.py
index c4543a1f..b2253f31 100644
--- a/test/test_remove_assert.py
+++ b/test/test_remove_assert.py
@@ -1,5 +1,6 @@
 import ast
 
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import compare_ast
 from python_minifier.rename import add_namespace, bind_names, resolve_names
 from python_minifier.transforms.remove_asserts import RemoveAsserts
@@ -8,6 +9,7 @@
 def remove_asserts(source):
     module = ast.parse(source, 'remove_asserts')
 
+    add_parent(module)
     add_namespace(module)
     bind_names(module)
     resolve_names(module)
diff --git a/test/test_remove_debug.py b/test/test_remove_debug.py
index f15fc467..c02d8e97 100644
--- a/test/test_remove_debug.py
+++ b/test/test_remove_debug.py
@@ -2,6 +2,7 @@
 
 import pytest
 
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import compare_ast
 from python_minifier.rename import add_namespace, bind_names, resolve_names
 from python_minifier.transforms.remove_debug import RemoveDebug
@@ -10,6 +11,7 @@
 def remove_debug(source):
     module = ast.parse(source, 'remove_debug')
 
+    add_parent(module)
     add_namespace(module)
     bind_names(module)
     resolve_names(module)
diff --git a/test/test_remove_exception_brackets.py b/test/test_remove_exception_brackets.py
index 0a81ca55..aecb55c5 100644
--- a/test/test_remove_exception_brackets.py
+++ b/test/test_remove_exception_brackets.py
@@ -3,6 +3,7 @@
 
 import pytest
 
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import compare_ast
 from python_minifier.rename import add_namespace, bind_names, resolve_names
 from python_minifier.transforms.remove_exception_brackets import remove_no_arg_exception_call
@@ -11,6 +12,7 @@
 def remove_brackets(source):
     module = ast.parse(source, 'remove_brackets')
 
+    add_parent(module)
     add_namespace(module)
     bind_names(module)
     resolve_names(module)
diff --git a/test/test_remove_literal_statements.py b/test/test_remove_literal_statements.py
index d66c0417..b9289fbd 100644
--- a/test/test_remove_literal_statements.py
+++ b/test/test_remove_literal_statements.py
@@ -1,5 +1,6 @@
 import ast
 
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import compare_ast
 from python_minifier.rename import add_namespace, bind_names, resolve_names
 from python_minifier.transforms.remove_literal_statements import RemoveLiteralStatements
@@ -8,6 +9,7 @@
 def remove_literals(source):
     module = ast.parse(source, 'test_remove_literal_statements')
 
+    add_parent(module)
     add_namespace(module)
     bind_names(module)
     resolve_names(module)
diff --git a/test/test_remove_pass.py b/test/test_remove_pass.py
index ba74af83..4947ad0e 100644
--- a/test/test_remove_pass.py
+++ b/test/test_remove_pass.py
@@ -1,5 +1,6 @@
 import ast
 
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import compare_ast
 from python_minifier.rename import add_namespace, bind_names, resolve_names
 from python_minifier.transforms.remove_pass import RemovePass
@@ -8,6 +9,7 @@
 def remove_literals(source):
     module = ast.parse(source, 'remove_literals')
 
+    add_parent(module)
     add_namespace(module)
     bind_names(module)
     resolve_names(module)
diff --git a/test/test_rename_builtins.py b/test/test_rename_builtins.py
index c54224f6..4fd7c3c9 100644
--- a/test/test_rename_builtins.py
+++ b/test/test_rename_builtins.py
@@ -7,6 +7,7 @@
 import ast
 
 from python_minifier import unparse
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import CompareError, compare_ast
 from python_minifier.rename import add_namespace, allow_rename_globals, allow_rename_locals, bind_names, rename, resolve_names
 
@@ -14,6 +15,7 @@
 def do_rename(source):
     # This will raise if the source file can't be parsed
     module = ast.parse(source, 'test_rename_bultins')
+    add_parent(module)
     add_namespace(module)
     bind_names(module)
     resolve_names(module)
diff --git a/test/test_rename_locals.py b/test/test_rename_locals.py
index 4529654a..ff2a16ff 100644
--- a/test/test_rename_locals.py
+++ b/test/test_rename_locals.py
@@ -10,6 +10,7 @@
 import pytest
 
 from python_minifier import unparse
+from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import CompareError, compare_ast
 from python_minifier.rename import add_namespace, allow_rename_globals, allow_rename_locals, bind_names, rename, resolve_names
 
@@ -18,6 +19,7 @@ def rename_locals(source):
 
     # This will raise if the source file can't be parsed
     module = ast.parse(source, 'test_rename_locals')
+    add_parent(module)
     add_namespace(module)
     bind_names(module)
     resolve_names(module)