Python lib2to3/libfuturize fixer for dict methods

fix_dict_methods.py

"""Fixer for some dict methods.

Imports these compat methods from `future.utils`

d.iterkeys() -> iterkeys(d)
d.iteritems() -> iteritems(d)
d.itervalues() -> itervalues(d)

d.viewkeys() -> viewkeys(d)
d.viewitems() -> viewitems(d)
d.viewvalues() -> viewvalues(d)

`futurize -w -n -f <package>.lib.fixes.fix_dict_methods .`
"""

from lib2to3 import pytree
from lib2to3 import fixer_base
from lib2to3.fixer_util import Call, Name
from libfuturize.fixer_util import touch_import_top


class FixDictMethods(fixer_base.BaseFix):
    BM_compatible = True

    PATTERN = """
    power< head=any+
         trailer< '.' method=('iterkeys'|'iteritems'|'itervalues'|
                              'viewkeys'|'viewitems'|'viewvalues') >
         parens=trailer< '(' ')' >
         tail=any*
    >
    """

    def transform(self, node, results):
        node_type = self.syms.power
        head = results["head"]
        method = results["method"][0]  # Extract node for method name
        tail = results["tail"]
        method_name = method.value

        touch_import_top(u'future.utils', method_name, node)

        head = [n.clone() for n in head]
        tail = [n.clone() for n in tail]
        head[0].prefix = ''
        return pytree.Node(
            node_type, [Call(Name(method_name), args=head, prefix=node.prefix)] + tail
        )

Might be possible to simplify this by removing tail from the PATTERN and return statement, but this works fine for me 🤷‍♂️