Implementing a selections and other things

autoload/textobj/haskell.vim

@@ -4,14 +4,17 @@ if !has('python')
 endif
 
 function! textobj#haskell#select_i()
-    if g:haskell_textobj_include_types == 0
-        python selectHaskellBinding(vim.current.buffer, vim.current.window.cursor[0], False)
-    else
-        python selectHaskellBinding(vim.current.buffer, vim.current.window.cursor[0], True)
-    endif
+    python select_haskell_block(vim.current.buffer, vim.current.window.cursor[0], False)
 
     let start_position = g:haskell_textobj_ret[0]
     let end_position = g:haskell_textobj_ret[1]
     return ['v', start_position, end_position]
 endfunction
 
+function! textobj#haskell#select_a()
+    python select_haskell_block(vim.current.buffer, vim.current.window.cursor[0], True)
+
+    let start_position = g:haskell_textobj_ret[0]
+    let end_position = g:haskell_textobj_ret[1]
+    return ['v', start_position, end_position]
+endfunction

plugin/textobj/haskell.vim

@@ -11,16 +11,13 @@ if !exists('g:haskell_textobj_path')
     endif
 endif
 
-if !exists('g:haskell_textobj_include_types')
-    let g:haskell_textobj_include_types = 0
-endif
-
 python import vim
 execute 'pyfile ' . g:haskell_textobj_path
 
 call textobj#user#plugin('haskell', {
       \ '-': {
       \     'select-i': 'ih', '*select-i-function*': 'textobj#haskell#select_i',
+      \     'select-a': 'ah', '*select-a-function*': 'textobj#haskell#select_a',
       \ },
     \})
 

python/haskell-textobj.py

@@ -1,46 +1,242 @@
 #!/usr/bin/python
-import vim
 
-def isTopBinding(text):
-    return False if len(text) == 0 or (text.startswith("--") or text[0].isspace()) else True
+VIM_RETURN_VAR = 'haskell_textobj_ret'
 
-def isStatement(text):
-    return True if len(text) > 0 and (text[0].isspace() and len(text.strip()) > 0) else False
+import sys
 
-def isTypeSignature(text):
-    words = text.strip().split(" ")
-    return True if len(words) > 3 and words[1] == "::" else False
+try:
+    import vim
+except ImportError:
+    print "Warning: Not running inside Vim."
 
-def find(content, index, cmpF, iterF):
-    i = index
-    while (i >= 0 and i < len(content)):
-        if cmpF(content[i]):
-            return (True, i)
-        i = iterF(i)
-    return (False, i)
 
-def setRetValue(start, end,  lines):
-    startPos = [0, start+1, 1, 0]
-    endPos = [0, end+1, len(lines[end]), 0]
-    vim.command("let g:haskell_textobj_ret="+str([startPos, endPos]))
+def vim_return(start_line, end_line, lines):
+    """
+    Return the selection extent to Vim by setting a flag variable.
+    The format returned is described in the Vim documentation for getpos().
+    """
+    buf_num = 0
+    offset = 0
+    start_col = 1
+    start_pos = [buf_num, start_line + 1, start_col, offset]
+
+    end_col = len(lines[end_line])
+    end_pos = [buf_num, end_line + 1, end_col, offset]
+
+    cmd = "let g:%s=%s" % (VIM_RETURN_VAR, str([start_pos, end_pos]))
+    vim.command(cmd)
+
+
+def select_haskell_block(lines, cursor, around):
+    """
+    Find the start and end location of the current haskell text object.
 
-def selectHaskellBinding(lines, cursor, includeType):
+
+    Arguments:
+        - lines: list of haskell source lines
+        - cursor: line index
+        - around: include auxiliary blocks?
+
+    When `around` is true, the return value will select more. Specifically,
+    it will:
+        - Select all import statements in a block.
+        - Select all clauses of a function as well as the type signature.
     """
-    Extract function binding from index in content
-    content: [String] list of haskell source lines
-    cursor:  zero-based line index
-    return:  [String] top level binding index resides in
+    start_line, end_line = find_block(lines, cursor - 1)
+    if around:
+        # Take care of expanding imports.
+        if is_import(start_line, end_line, lines):
+            new = extend_imports(start_line, end_line, lines)
+            while new is not None:
+                start_line, end_line = new
+                new = extend_imports(start_line, end_line, lines)
+
+        # Take care of expanding pattern matches.
+        if is_decl(start_line, end_line, lines):
+            new = extend_decls(start_line, end_line, lines)
+            while new is not None:
+                start_line, end_line = new
+                new = extend_decls(start_line, end_line, lines)
+
+        # Add a type signature if necessary.
+        if is_decl(start_line, end_line, lines):
+            start_line, end_line = extend_typesig(start_line, end_line, lines)
+    vim_return(start_line, end_line, lines)
+
+
+def extend_decls(start_line, end_line, lines):
+    if start_line - 1 >= 0:
+        start2, end2 = find_block(lines, start_line - 1)
+        if start2 != start_line and is_decl(start2, end2, lines):
+            if lines[start_line].startswith(lines[start2].split()[0]):
+                return start2, end_line
+
+    if end_line + 1 <= len(lines) - 1:
+        start3, end3 = find_block(lines, end_line + 1)
+        if end3 != end_line and is_decl(start3, end3, lines):
+            if lines[start_line].startswith(lines[start3].split()[0]):
+                return start_line, end3
+
+    # No more cases
+    return None
+
+
+def extend_imports(start_line, end_line, lines):
+    if start_line - 1 >= 0:
+        start2, end2 = find_block(lines, start_line - 1)
+        if start2 != start_line and is_import(start2, end2, lines):
+            return start2, end_line
+
+    if end_line + 1 <= len(lines) - 1:
+        start3, end3 = find_block(lines, end_line + 1)
+        if end3 != end_line and is_import(start3, end3, lines):
+            return start_line, end3
+
+    # Signal no more imports to add
+    return None
+
+
+def is_import(start, end, lines):
+    return lines[start].strip().startswith("import")
+
+
+def is_decl(start, end, lines):
+    line = lines[start].strip()
+    decl = "=" in line
+    if not decl:
+        return False
+
+    if any(tok in line for tok in ["data", "newtype"]):
+        return False
+
+    return True
+
+
+def extend_typesig(start_line, end_line, lines):
+    start2, end2 = find_block(lines, start_line - 1)
+
+    first_line = start2
+    while is_comment(lines[first_line]):
+        first_line += 1
+
+    if "::" in lines[first_line].split()[1]:
+        return start2, end_line
+    else:
+        return start_line, end_line
+
+
+def indent_level(line):
+    """Return the indent level of the line (measured in spaces)"""
+    # Make sure it has something on it.
+    if not line.strip():
+        raise ValueError("Empty line has no indent level")
+
+    level = 0
+    for i, char in enumerate(line):
+        if char != ' ' and char != '\t':
+            return level
+        elif char == ' ':
+            level += 1
+        elif char == '\t':
+            level += 8
+
+
+def empty(line):
+    return not line.strip()
+
+
+def indented(line):
+    return indent_level(line) > 0
+
+
+def has_start_block(line):
+    """Whether this line is the beginning of a block."""
+    maybe_type = len(line.split("::")) == 2
+    if maybe_type:
+        before, after = line.split("::")
+        if before.count(" ") <= 0:
+            return True
+    return any(line.startswith(start_token)
+               for start_token in ["data", "newtype", "import"])
+
+
+def is_comment(line):
+    line = line.strip()
+    is_pragma = line.startswith("{-#")
+    return line.startswith("--") or (line.startswith("{-") and not is_pragma)
+
+
+def find_block(lines, index):
+    """Find a block that the cursor is in.
+
+    Arguments:
+        - lines: all the lines in the file.
+        - index: the line number of the cursor.
+        - around: Whether to include surrounding blocks.
     """
-    backward = lambda x : x - 1
-    forward  = lambda x : x + 1
-    index = cursor - 1
+    # Move the cursor until we find a non-empty line.
+    while index < len(lines) and empty(lines[index]):
+        index += 1
+
+    # Start by only including the line we're on.
+    start_index = index
+    end_index = index
+
+    # Expand the selection upwards.
+    def expand_upwards():
+        # Can't go past file start.
+        if start_index == 0:
+            return False
+
+        current_line = lines[start_index]
+        previous_line = lines[start_index - 1]
+
+        if is_comment(previous_line):
+            return True
+
+        # If this is the start of a block, go no further.
+        if not empty(current_line) and not indented(current_line):
+            return False
+
+        return True
+
+    while expand_upwards():
+        start_index -= 1
+
+    def expand_downwards():
+        # Can't go past file end.
+        if end_index == len(lines) - 1:
+            return False
+
+        current_line = lines[end_index]
+        next_line = lines[end_index + 1]
+
+        if is_comment(current_line) and is_comment(next_line):
+            return True
+
+        # Can't encroach on the next start block.
+        if has_start_block(next_line):
+            return False
+
+        return empty(next_line) or indented(next_line)
+
+    while expand_downwards():
+        end_index += 1
 
-    found, bStart = find(lines, index, isTopBinding, backward)
-    found, bNext  = find(lines, index+1, isTopBinding, forward)
-    found, bEnd   = find(lines, bNext, isStatement, backward) if found else (True, bNext-1)
-    bEnd = bStart if bEnd < bStart else bEnd
+    # Trim the selection to avoid newlines.
+    while empty(lines[start_index]):
+        start_index += 1
+    while empty(lines[end_index]):
+        end_index -= 1
 
-    if includeType and bStart > 0 and isTypeSignature(lines[bStart-1]):
-        bStart = bStart - 1
+    return start_index, end_index
 
-    setRetValue(bStart, bEnd, lines)
+if __name__ == "__main__" and 'vim' not in sys.modules:
+    lines = open(sys.argv[1]).readlines()
+    for ind in xrange(len(lines)):
+        start, end = find_block(lines, ind)
+        for i in xrange(start, end + 1):
+            print lines[i][:-1]
+        print ind, find_block(lines, ind)
+        raw_input()
+        print '---'