fix_dict.py

"""Fixer for six.iterkeys() -> .keys(), and similarly for iteritems and itervalues."""

from lib2to3 import fixer_util, fixer_base
from lib2to3 import pytree
from lib2to3.fixes import fix_dict
from lib2to3.fixer_util import Call, Name
import libpostmodernize


class FixDict(fix_dict.FixDict):
    PATTERN = """
    power< six='six' trailer< '.' method=('iterkeys'|'iteritems'|'itervalues'|
                                      'viewkeys'|'viewitems'|'viewvalues') >
            trailer< '(' arg=any ')' > >
    |
    power< method=('iterkeys'|'iteritems'|'itervalues'|
                   'viewkeys'|'viewitems'|'viewvalues') trailer< '(' arg=any ')' > >
    """

    def transform(self, node, results):
        syms = self.syms

        in_special_context = self.in_special_context(node)
        name = results['method'][0]
        arg = results['arg']
        # XXX we discard arg.prefix which can lose a comment in some situations
        arg.prefix = node.prefix
        suffix = arg.get_suffix()
        arg.remove()
        new_node = pytree.Node(syms.power, [arg, Name("." + name.value[4:] + "()")])

        if name.value.startswith("iter") and not in_special_context:
            pre = new_node.prefix
            new_node.prefix = ""
            new_node = Call(Name("iter", prefix=pre), [new_node])

        node.replace(new_node)
        # WHAT DO IF NOT next_sibling???
        sib = new_node.next_sibling
        if sib:
            sib.prefix = suffix.rstrip() + sib.prefix



    def in_special_context(self, node):
        # Redefined from parent class to make "for x in d.items()" count as
        # in special context; 2to3 only counts for loops as special context
        # for the iter* methods.
        return super().in_special_context(node, True)