"""
this file has been tested in python 3.8
"""

import collections
from io import StringIO
from lib2to3.patcomp import PatternCompiler
from lib2to3 import pytree, fixer_base, refactor, pgen2
from lib2to3.pygram import python_symbols, python_grammar
from lib2to3.pytree import Node
from lib2to3.fixer_util import Name, Dot

def create_node(module_path):
    dotted_parts = module_path.split('.')
    node = Node(python_symbols.power, [])
    children = [Name(dotted_parts[0]),]
    for part in dotted_parts[1:]:
        children.append(Node(python_symbols.trailer, [Dot(), Name(part)]))
    node.children = children
    return node


class Fixer1(fixer_base.BaseFix):
    PATTERN = """
              'MyClass'
              """
    def transform(self, node, results):
        """
        convert `print(type(MyClass))` to `print(type(path1.to1.MyClass))`
        """

        return create_node('path1.to1.MyClass')


class Fixer2(fixer_base.BaseFix):
    PATTERN = """
              power< 'path1' trailer< '.' 'to1' > trailer< '.' 'MyClass' > >
              """
    def transform(self, node, results):
        """
        convert `print(type(path1.to1.MyClass))` to `print(type(path2.to2.MyClass))`
        """
        node.children[0].value = 'path2'
        node.children[1].children[1].value = 'to2'

source = """
print(type(MyClass))
"""
driver = pgen2.driver.Driver(python_grammar, convert=pytree.convert)
tree = driver.parse_stream(StringIO(source.strip() + "\n"))

fixer1 = Fixer1({}, [])
fixer2 = Fixer2({}, [])
bmi_post_order = [fixer1, fixer2]


def _default_wrong_behavior():
    """
    the default behavior which is wrong
    """
    bmi_post_order_heads = refactor._get_headnode_dict(bmi_post_order)
    tool = refactor.RefactoringTool([], {}, explicit=True)
    tool.traverse_by(bmi_post_order_heads, tree.post_order())
    print(tree) # output is "print(type(path1.to1.MyClass))", which should be "print(type(path2.to2.MyClass))"
    print(repr(tree))

def _correct_behavior():
    """
    run this code after applying patch in https://github.com/python/cpython/pull/21988
    """
    accept_type_dict = collections.defaultdict(set)
    refactor._get_accept_type_dict(bmi_post_order, accept_type_dict)
    tool = refactor.RefactoringTool([], {}, explicit=True)
    tool.fixer_accept_type = accept_type_dict
    tool.traverse_by(bmi_post_order, tree)
    print(tree) # output is "print(type(path2.to2.MyClass))", which is correct
    print(repr(tree))

_default_wrong_behavior()
# comment the prior line and uncomment the following line to get the correct behavior
#_correct_behavior()