import_sorter.py

import distutils.sysconfig as sysconfig
import re

import sublime
import sublime_plugin

FROM_IMPORT_RE = re.compile(r'^\s*from (?P<package>\.?\w*)(?P<package_extra>\S*)\s+import\s+(?P<subpackages>.+)')
PLAIN_IMPORT_RE = re.compile(r'^\s*import (?P<package>\.?\w*)(?P<package_extra>\S*)')
CONTINUATION_RE = re.compile(r'.*(\\\s*)$')
PARENTHESISED_START_RE = re.compile(r'^.+? import \(')
PARENTHESISED_END_RE = re.compile(r'.*?\)\s*$')

class PythonImportSorterCommand(sublime_plugin.TextCommand):
    def run(self, edit):
        local_package_prefixes = self.view.settings().get('my_py_packages', [])
        
        # Build a list of standard library modules
        stdlib_path = sysconfig.get_python_lib(standard_lib=True)
        
        for selection_range in self.view.sel():
            selection = self.view.substr(selection_range)
            imports = {
                # First group is import, second is from...import
                'stdlib': ([], []),
                'thirdparty': ([], []),
                'local': ([], []),
                'relative': ([], []),
            }
            lines = [l.rstrip() for l in selection.splitlines()]
            in_parenthesis = False
            
            for line_index, line in enumerate(lines):
                # If the line is empty, don't bother with it
                if not line.strip():
                    continue
                
                # If the line ends in a continuation, keep consuming until we have
                # the entire line
                continuation_match = CONTINUATION_RE.match(line)
                while continuation_match is not None:
                    line = line[:continuation_match.start(1)].rstrip()
                    continuation_line = lines[line_index + 1]
                    line = '%s %s' % (line, continuation_line.lstrip())
                    continuation_match = CONTINUATION_RE.match(line)
                    # Skip the next line
                    lines[line_index + 1] = ''
                    line_index += 1
                
                # Similarly, if we are in a parenthasis, keep consuming until the end
                in_parenthesis = in_parenthesis or PARENTHESISED_START_RE.match(line)
                while in_parenthesis:
                    continuation_line = lines[line_index + 1]
                    lines[line_index + 1] = ''
                    line = '%s %s' % (line.rstrip(), continuation_line.lstrip())
                    in_parenthesis = not PARENTHESISED_END_RE.match(line)
                    if not in_parenthesis:
                        # We are at the end of the parenthesised line, remove the actual parens
                        line = re.sub(r'[\(\)]', '', line)
                    line_index += 1
                
                # Determine whether this is a from...import or just an import
                if PLAIN_IMPORT_RE.match(line):
                    match = PLAIN_IMPORT_RE.match(line)
                    line = {
                        'package': match.group('package'),
                        'fullpackage': match.group('package') + match.group('package_extra'),
                        'subpackages': (),
                        'raw': line,
                        'match': match
                    }
                    destination = 0
                elif FROM_IMPORT_RE.match(line):
                    match = FROM_IMPORT_RE.match(line)
                    line = {
                        'package': match.group('package'),
                        'fullpackage': match.group('package') + match.group('package_extra'),
                        'subpackages': sorted([i.strip() for i in re.split(r',\s*', match.group(3)) if i.strip()]),
                        'raw': line,
                        'match': match
                    }
                    destination = 1
                else:
                    # Don't do anything
                    assert 0, repr(line)
                
                # Process subpackages as appropriate
                if line['subpackages']:
                    span = line['match'].span('subpackages')
                    raw = '%s%s%s' % (line['raw'][:span[0]], ', '.join(line['subpackages']), line['raw'][span[1]:])
                    line['raw'] = raw
                
                # Add to the appropriate list
                if line['package'].split('.', 1)[0] in local_package_prefixes:
                    # A local package
                    key = 'local'
                elif line['package'].startswith('.'):
                    # A relative import
                    key = 'relative'
                else:
                    try:
                        module = __import__(line['package'])
                        assert module.__file__.startswith(stdlib_path)
                    except (ImportError, AssertionError):
                        key = 'thirdparty'
                    except AttributeError:
                        # When __file__ is not present, it must be a builtin
                        key = 'stdlib'
                    else:
                        key = 'stdlib'
                
                imports[key][destination].append(line)
            
            # Now we have parsed the list of imports, assemble them into a sensible order
            order = ('stdlib', 'thirdparty', 'local', 'relative')
            output = [] # Import groups
            package_comparator = lambda x, y: cmp(x['raw'], y['raw'])
            
            for key in order:
                plain_lines, from_lines = imports[key]
                # Don't do anything if there are no matches for that style of import
                if not (plain_lines or from_lines):
                    continue
                
                out = []
                out.extend([i['raw'] for i in sorted(plain_lines, package_comparator)])
                out.extend([i['raw'] for i in sorted(from_lines, package_comparator)])
                output.append('\n'.join(out))
            
            output = '\n\n'.join(output)
            
            # Take care of line wrapping at the appropriate column
            try:
                wrap_column = self.view.settings().get('rulers')[0]
            except (IndexError, TypeError):
                # No wrap column is defined
                print 'passing'
                pass
            else:
                output = output.splitlines()
                for i, line in enumerate(output):
                    if len(line) > wrap_column:
                        head, tail = line[:(wrap_column - 2)], line[(wrap_column - 2):]
                        
                        # Try to sensibly wrap, by searching for the last bit of whitespace in the head and moving it
                        # to the tail
                        try:
                            match = list(re.finditer(r'\s+', head))[-1]
                        except IndexError:
                            continue
                        
                        if match:
                            span = match.span()
                            tail = '%s%s' % (head[span[0]:], tail.lstrip())
                            head = head[:span[0]].rstrip()
                            
                            head = '%s \\' % head.rstrip()
                            tail = '    %s' % tail.lstrip()
                            output[i] = head
                            output.insert(i + 1, tail)
                
                output = '\n'.join(output)
            
            self.view.replace(edit, selection_range, output)