""" Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """



# Local imports

from .. import fixer_base

from ..fixer_util import BlankLine



class FixItertoolsImports(fixer_base.BaseFix):

    PATTERN = """

              import_from< 'from' 'itertools' 'import' imports=any >

              """ %(locals())



    def transform(self, node, results):

        imports = results['imports']

        children = imports.children[:] or [imports]

        for child in children:

            if not hasattr(child, 'value'):

                # Handle 'import ... as ...'

                continue

            if child.value in ('imap', 'izip', 'ifilter'):

                # The value must be set to none in case child == import,

                # so that the test for empty imports will work out

                child.value = None

                child.remove()

            elif child.value == 'ifilterfalse':

                node.changed()

                child.value = 'filterfalse'



        # Make sure the import statement is still sane

        children = imports.children[:] or [imports]

        remove_comma = True

        for child in children:

            if remove_comma and getattr(child, 'value', None) == ',':

                child.remove()

            else:

                remove_comma ^= True



        if unicode(children[-1]) == ',':

            children[-1].remove()



        # If there are no imports left, just get rid of the entire statement

        if not (imports.children or getattr(imports, 'value', None)):

            p = node.get_prefix()

            node = BlankLine()

            node.prefix = p

        return node

