Misago/extras/fixdictsformatting.py

from __future__ import unicode_literals

import sys

from lib2to3.pytree import Node, Leaf
from lib2to3.fixer_util import token, syms

from yapf.yapflib import pytree_utils

from django.utils import six

from .config import yapf as yapf_config


MAX_LINE_LENGTH = yapf_config.getint('style', 'column_limit') + 1


def fix_formatting(filesource):
    if not ('{'  in filesource and ('[' in filesource or '(' in filesource)):
        return filesource

    tree = pytree_utils.ParseCodeToTree(filesource)
    for node in tree.children:
        walk_tree(node, node.children)
    return six.text_type(tree)


def walk_tree(node, children):
    for item in children:
        if item.type == syms.dictsetmaker:
            indent = item.parent.children[-1].column
            walk_dict_tree(item, item.children, indent)
        else:
            walk_tree(item, item.children)


def walk_dict_tree(node, children, indent):
    for item in children:
        prev = item.prev_sibling
        if isinstance(prev, Leaf) and prev.value == ':':
            if isinstance(item, Leaf):
                if six.text_type(item).startswith("\n"):
                    item.replace(Leaf(
                        item.type,
                        item.value,
                        prefix=' ',
                    ))
            elif six.text_type(item).strip()[0] in ('[', '{'):
                walk_tree(item, item.children)
            else:
                walk_dedent_tree(item, item.children, indent)


def walk_dedent_tree(node, children, indent):
    force_split_next = False
    for item in children:
        prev = item.prev_sibling
        if not prev:
            if isinstance(item, Leaf) and six.text_type(item).startswith("\n"):
                prev = node.prev_sibling
                next = node.next_sibling
                final_length = 0

                if prev and "\n" not in six.text_type(node).strip():
                    final_length = prev.column + len(six.text_type(node).strip()) + 3

                item.replace(Leaf(
                    item.type,
                    item.value,
                    prefix=' ',
                ))

                if final_length and final_length > MAX_LINE_LENGTH:
                    # tell next call to walk_dedent_tree_node that we need
                    # different stringformat tactic
                    force_split_next = True
        elif isinstance(item, Node):
            if node.type == syms.power:
                for subitem in item.children[1:]:
                    walk_dedent_power_node(subitem, subitem.children, indent)
            else:
                for subitem in item.children[1:]:
                    walk_dedent_tree_node(subitem, subitem.children, indent, force_split_next)
                    force_split_next = False


def walk_dedent_tree_node(node, children, indent, force_split_next=False):
    if six.text_type(node).startswith("\n"):
        if isinstance(node, Leaf):
            prev = node.prev_sibling
            next = node.next_sibling

            is_followup = prev and prev.type == token.STRING and node.type == token.STRING
            if is_followup:
                new_value = node.value

                # insert linebreak after last string in braces, so its closing brace moves to new line
                if not node.next_sibling:
                    closing_bracket = node.parent.parent.children[-1]
                    if not six.text_type(closing_bracket).startswith("\n"):
                        new_value = "%s\n%s" % (node.value, (' ' * (indent + 4)))

                node.replace(Leaf(
                    node.type,
                    new_value,
                    prefix="\n%s" % (' ' * (indent + 8)),
                ))
            else:
                if six.text_type(node).strip() in (')', '}'):
                    new_prefix = "\n%s" % (' ' * (indent + 4))
                else:
                    new_prefix = "\n%s" % (' ' * (indent + 8))

                node.replace(Leaf(
                    node.type,
                    node.value,
                    prefix=new_prefix
                ))
        else:
            for item in children:
                walk_dedent_tree_node(item, item.children, indent)
    elif isinstance(node, Leaf):
        if node.type == token.STRING:
            strings_tuple = node.parent.parent

            prev = node.prev_sibling
            next = node.next_sibling

            is_opening = prev is None and six.text_type(strings_tuple).strip()[0] == '('
            has_followup = next and next.type == token.STRING

            if is_opening and has_followup:
                node.replace(Leaf(
                    node.type,
                    node.value,
                    prefix="\n%s" % (' ' * (indent + 8)),
                ))
            elif force_split_next:
                node.replace(Leaf(
                    node.type,
                    "%s\n%s" % (node.value, (' ' * (indent + 4))),
                    prefix="\n%s" % (' ' * (indent + 8)),
                ))
    else:
        for item in children:
            walk_dedent_tree_node(item, item.children, indent)


def walk_dedent_power_node(node, children, indent):
    if isinstance(node, Leaf):
        if six.text_type(node).startswith("\n"):
            node.replace(Leaf(
                node.type,
                node.value,
                prefix=node.prefix[:-4],
            ))
    else:
        for item in children:
            walk_dedent_power_node(item, item.children, indent)