From 2983d4478fa03dbbd495ff30e86320537b9dd0a1 Mon Sep 17 00:00:00 2001
From: Anthony Sottile <asottile@umich.edu>
Date: Tue, 10 Feb 2015 08:14:53 -0800
Subject: [PATCH] Use the tokenizer for great success

---
 pre_commit_hooks/string_fixer.py | 52 +++++++++++++++------
 tests/string_fixer_test.py       | 78 ++++++++++----------------------
 2 files changed, 63 insertions(+), 67 deletions(-)

diff --git a/pre_commit_hooks/string_fixer.py b/pre_commit_hooks/string_fixer.py
index 9d2213a..9ef7a37 100644
--- a/pre_commit_hooks/string_fixer.py
+++ b/pre_commit_hooks/string_fixer.py
@@ -3,34 +3,60 @@ from __future__ import print_function
 from __future__ import unicode_literals
 
 import argparse
-import re
+import io
 import tokenize
 
 
 double_quote_starts = tuple(s for s in tokenize.single_quoted if '"' in s)
-compiled_tokenize_string = re.compile('(?<!")' + tokenize.String + '(?!")')
 
 
-def handle_match(m):
-    string = m.group(0)
+def handle_match(token_text):
+    if '"""' in token_text or "'''" in token_text:
+        return token_text
 
     for double_quote_start in double_quote_starts:
-        if string.startswith(double_quote_start):
-            meat = string[len(double_quote_start):-1]
+        if token_text.startswith(double_quote_start):
+            meat = token_text[len(double_quote_start):-1]
             if '"' in meat or "'" in meat:
                 break
             return double_quote_start.replace('"', "'") + meat + "'"
-    return string
+    return token_text
+
+
+def get_line_offsets_by_line_no(src):
+    # Padded so we can index with line number
+    offsets = [None, 0]
+    for line in src.splitlines():
+        offsets.append(offsets[-1] + len(line) + 1)
+    return offsets
 
 
 def fix_strings(filename):
-    contents = open(filename).read()
-    new_contents = compiled_tokenize_string.sub(handle_match, contents)
-    retval = int(new_contents != contents)
-    if retval:
-        with open(filename, 'w') as write_handle:
+    contents = io.open(filename).read()
+    line_offsets = get_line_offsets_by_line_no(contents)
+
+    # Basically a mutable string
+    splitcontents = list(contents)
+
+    # Iterate in reverse so the offsets are always correct
+    tokens = reversed(list(tokenize.generate_tokens(
+        io.StringIO(contents).readline,
+    )))
+    for token_type, token_text, (srow, scol), (erow, ecol), _ in tokens:
+        if token_type == tokenize.STRING:
+            new_text = handle_match(token_text)
+            splitcontents[
+                line_offsets[srow] + scol:
+                line_offsets[erow] + ecol
+            ] = new_text
+
+    new_contents = ''.join(splitcontents)
+    if contents != new_contents:
+        with io.open(filename, 'w') as write_handle:
             write_handle.write(new_contents)
-    return retval
+        return 1
+    else:
+        return 0
 
 
 def main(argv=None):
diff --git a/tests/string_fixer_test.py b/tests/string_fixer_test.py
index 15b9f19..6305618 100644
--- a/tests/string_fixer_test.py
+++ b/tests/string_fixer_test.py
@@ -2,79 +2,49 @@ from __future__ import absolute_import
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import textwrap
+
 import pytest
 
 from pre_commit_hooks.string_fixer import main
 
 TESTS = (
     # Base cases
-    (
-        "''",
-        "''",
-        0
-    ),
-    (
-        '""',
-        "''",
-        1
-    ),
-    (
-        r'"\'"',
-        r'"\'"',
-        0
-    ),
-    (
-        r'"\""',
-        r'"\""',
-        0
-    ),
-    (
-        r"'\"\"'",
-        r"'\"\"'",
-        0
-    ),
+    ("''", "''", 0),
+    ('""', "''", 1),
+    (r'"\'"', r'"\'"', 0),
+    (r'"\""', r'"\""', 0),
+    (r"'\"\"'", r"'\"\"'", 0),
     # String somewhere in the line
-    (
-        'x = "foo"',
-        "x = 'foo'",
-        1
-    ),
+    ('x = "foo"', "x = 'foo'", 1),
     # Test escaped characters
-    (
-        r'"\'"',
-        r'"\'"',
-        0
-    ),
+    (r'"\'"', r'"\'"', 0),
     # Docstring
+    ('""" Foo """', '""" Foo """', 0),
     (
-        '""" Foo """',
-        '""" Foo """',
-        0
-    ),
-    # Fuck it, won't even try to fix
-    (
-        """
-        x = " \\n
-        foo \\n
+        textwrap.dedent("""
+        x = " \\
+        foo \\
         "\n
-        """,
-        """
-        x = " \\n
-        foo \\n
-        "\n
-        """,
-        0
+        """),
+        textwrap.dedent("""
+        x = ' \\
+        foo \\
+        '\n
+        """),
+        1,
     ),
+    ('"foo""bar"', "'foo''bar'", 1),
 )
 
 
-@pytest.mark.parametrize(('input_s', 'expected_output', 'expected_retval'), TESTS)
-def test_rewrite(input_s, expected_output, expected_retval, tmpdir):
+@pytest.mark.parametrize(('input_s', 'output', 'expected_retval'), TESTS)
+def test_rewrite(input_s, output, expected_retval, tmpdir):
     tmpfile = tmpdir.join('file.txt')
 
     with open(tmpfile.strpath, 'w') as f:
         f.write(input_s)
 
     retval = main([tmpfile.strpath])
-    assert tmpfile.read() == expected_output
+    assert tmpfile.read() == output
     assert retval == expected_retval