Merge pull request #1643 from MarcoGorelli/negate-pygrep

asottile · web-flow · commit 01f1a0090e52 · 2020-10-17T11:14:19.000-07:00
ENH add --negate flag to pygrep
diff --git a/pre_commit/languages/pygrep.py b/pre_commit/languages/pygrep.py
@@ -1,6 +1,7 @@
 import argparse
 import re
 import sys
+from typing import NamedTuple
 from typing import Optional
 from typing import Pattern
 from typing import Sequence
@@ -45,6 +46,46 @@ def _process_filename_at_once(pattern: Pattern[bytes], filename: str) -> int:
     return retv
 
 
+def _process_filename_by_line_negated(
+        pattern: Pattern[bytes],
+        filename: str,
+) -> int:
+    with open(filename, 'rb') as f:
+        for line in f:
+            if pattern.search(line):
+                return 0
+        else:
+            output.write_line(filename)
+            return 1
+
+
+def _process_filename_at_once_negated(
+        pattern: Pattern[bytes],
+        filename: str,
+) -> int:
+    with open(filename, 'rb') as f:
+        contents = f.read()
+    match = pattern.search(contents)
+    if match:
+        return 0
+    else:
+        output.write_line(filename)
+        return 1
+
+
+class Choice(NamedTuple):
+    multiline: bool
+    negate: bool
+
+
+FNS = {
+    Choice(multiline=True, negate=True): _process_filename_at_once_negated,
+    Choice(multiline=True, negate=False): _process_filename_at_once,
+    Choice(multiline=False, negate=True): _process_filename_by_line_negated,
+    Choice(multiline=False, negate=False): _process_filename_by_line,
+}
+
+
 def run_hook(
         hook: Hook,
         file_args: Sequence[str],
@@ -64,6 +105,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
     )
     parser.add_argument('-i', '--ignore-case', action='store_true')
     parser.add_argument('--multiline', action='store_true')
+    parser.add_argument('--negate', action='store_true')
     parser.add_argument('pattern', help='python regex pattern.')
     parser.add_argument('filenames', nargs='*')
     args = parser.parse_args(argv)
@@ -75,11 +117,9 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
     pattern = re.compile(args.pattern.encode(), flags)
 
     retv = 0
+    process_fn = FNS[Choice(multiline=args.multiline, negate=args.negate)]
     for filename in args.filenames:
-        if args.multiline:
-            retv |= _process_filename_at_once(pattern, filename)
-        else:
-            retv |= _process_filename_by_line(pattern, filename)
+        retv |= process_fn(pattern, filename)
     return retv
 
 
diff --git a/tests/languages/pygrep_test.py b/tests/languages/pygrep_test.py
@@ -8,6 +8,9 @@ def some_files(tmpdir):
     tmpdir.join('f1').write_binary(b'foo\nbar\n')
     tmpdir.join('f2').write_binary(b'[INFO] hi\n')
     tmpdir.join('f3').write_binary(b"with'quotes\n")
+    tmpdir.join('f4').write_binary(b'foo\npattern\nbar\n')
+    tmpdir.join('f5').write_binary(b'[INFO] hi\npattern\nbar')
+    tmpdir.join('f6').write_binary(b"pattern\nbarwith'foo\n")
     with tmpdir.as_cwd():
         yield
 
@@ -30,6 +33,58 @@ def test_main(cap_out, pattern, expected_retcode, expected_out):
     assert out == expected_out
 
 
+@pytest.mark.usefixtures('some_files')
+def test_negate_by_line_no_match(cap_out):
+    ret = pygrep.main(('pattern\nbar', 'f4', 'f5', 'f6', '--negate'))
+    out = cap_out.get()
+    assert ret == 1
+    assert out == 'f4\nf5\nf6\n'
+
+
+@pytest.mark.usefixtures('some_files')
+def test_negate_by_line_two_match(cap_out):
+    ret = pygrep.main(('foo', 'f4', 'f5', 'f6', '--negate'))
+    out = cap_out.get()
+    assert ret == 1
+    assert out == 'f5\n'
+
+
+@pytest.mark.usefixtures('some_files')
+def test_negate_by_line_all_match(cap_out):
+    ret = pygrep.main(('pattern', 'f4', 'f5', 'f6', '--negate'))
+    out = cap_out.get()
+    assert ret == 0
+    assert out == ''
+
+
+@pytest.mark.usefixtures('some_files')
+def test_negate_by_file_no_match(cap_out):
+    ret = pygrep.main(('baz', 'f4', 'f5', 'f6', '--negate', '--multiline'))
+    out = cap_out.get()
+    assert ret == 1
+    assert out == 'f4\nf5\nf6\n'
+
+
+@pytest.mark.usefixtures('some_files')
+def test_negate_by_file_one_match(cap_out):
+    ret = pygrep.main(
+        ('foo\npattern', 'f4', 'f5', 'f6', '--negate', '--multiline'),
+    )
+    out = cap_out.get()
+    assert ret == 1
+    assert out == 'f5\nf6\n'
+
+
+@pytest.mark.usefixtures('some_files')
+def test_negate_by_file_all_match(cap_out):
+    ret = pygrep.main(
+        ('pattern\nbar', 'f4', 'f5', 'f6', '--negate', '--multiline'),
+    )
+    out = cap_out.get()
+    assert ret == 0
+    assert out == ''
+
+
 @pytest.mark.usefixtures('some_files')
 def test_ignore_case(cap_out):
     ret = pygrep.main(('--ignore-case', 'info', 'f1', 'f2', 'f3'))
