Vim: Use shorter textwidth in comments and docstrings

Question

From the mighty PEP 8:

[P]lease limit all lines to a maximum of 79 characters. For flowing long blocks of text (docstrings or comments), limiting the length to 72 characters is recommended.

When editing Python code in Vim, I set my textwidth to 79, and Vim automatically wraps long lines of Python code for me when I hit the character limit.

But in comments and docstrings, I need to wrap text at 72 characters instead. Is there any way to make Vim automatically set textwidth to 72 when I'm in a comment or docstring, and set it back when I'm done?

Answer 1

So, I've never done any Vim scripting before, but based on this question about doing something similar in C and this tip for checking if you're currently in a comment, I've hacked together a solution.

By default, this uses the PEP8-suggested widths of 79 characters for normal lines and 72 characters for comments, but you can override them by letting g:python_normal_text_width or g:python_comment_text_width variables, respectively. (Personally, I wrap normal lines at 78 characters.)

Drop this baby in your .vimrc and you should be good to go. I may package this up as a plugin later.

function! GetPythonTextWidth()
    if !exists('g:python_normal_text_width')
        let normal_text_width = 79
    else
        let normal_text_width = g:python_normal_text_width
    endif

    if !exists('g:python_comment_text_width')
        let comment_text_width = 72
    else
        let comment_text_width = g:python_comment_text_width
    endif

    let cur_syntax = synIDattr(synIDtrans(synID(line("."), col("."), 0)), "name")
    if cur_syntax == "Comment"
        return comment_text_width
    elseif cur_syntax == "String"
        " Check to see if we're in a docstring
        let lnum = line(".")
        while lnum >= 1 && synIDattr(synIDtrans(synID(lnum, col([lnum, "$"]) - 1, 0)), "name") == "String"
            if match(getline(lnum), "\\('''\\|\"\"\"\\)") > -1
                " Assume that any longstring is a docstring
                return comment_text_width
            endif
            let lnum -= 1
        endwhile
    endif

    return normal_text_width
endfunction

augroup pep8
    au!
    autocmd CursorMoved,CursorMovedI * :if &ft == 'python' | :exe 'setlocal textwidth='.GetPythonTextWidth() | :endif
augroup END