import pn
import scintilla
from pypn.decorators import script

def count_leading_spaces(input_str):
"""http://refactormycode.com/codes/1041-filter-by-indentation-level"""
    space_count = 0
    for char in input_str:
        if char.isspace():
            space_count += 1
        else:
            return space_count

@script("Align Columns", "Text")
def AlignColumns():
"""Align text into columns, preserving the leading whitespace.

This:

a b cc
aa b c
a bbb c

Becomes this:

a  b   cc
aa b   c
a  bbb c
"""
    editor = scintilla.Scintilla(pn.CurrentDoc())
    editor.BeginUndoAction()
    selection = editor.GetTextRange(editor.SelectionStart, editor.SelectionEnd)

    replace = align_cols(selection)
    editor.ReplaceSel(replace)
    editor.EndUndoAction()

def align_cols(text):
    leadingSpaces = 0
    maxColSize = []
    for line in text.splitlines(0):
        i = 0
        leadingSpaces = max(leadingSpaces, count_leading_spaces(line))
        for col in line.split():
            if (len(maxColSize) - 1 < i):
                maxColSize.append(len(col))
            else:
                maxColSize[i] = max(maxColSize[i], len(col))
            i = i + 1

    newLines = []
    
    for line in text.splitlines(0):
        cols = line.split();
        formatStr = ' ' * leadingSpaces + ' '.join(["%%-%ds" % val for val in maxColSize[:len(cols)]])
        newLines.append((formatStr % tuple(line.split())).rstrip())

    return '\r\n'.join(newLines)