[PATCH 9 of 9] grep: add support for max count

[Idan Kamara]

# HG changeset patch
# User Idan Kamara <idankk86 at gmail.com>
# Date 1350073314 -7200
# Node ID f01d8c64541fa2bfc3c085559510454f6e8e6019
# Parent  fc9b6a59bb64e8fff98671f2e6627f10bd889d0f
grep: add support for max count

diff --git a/mercurial/commands.py b/mercurial/commands.py
--- a/mercurial/commands.py
+++ b/mercurial/commands.py
@@ -2889,6 +2889,7 @@
      _('use this pattern to find matches (must be used if a pattern'
        ' starts with -), multiple patterns are or-ed'), _('PATTERN')),
     ('', 'invert-match', None, _('invert the sense of matching')),
+    ('m', 'max-count', '', _('stop matching after NUM matches'), _('NUM')),
     ('r', 'rev', [],
      _('only search files changed within revision range'), _('REV')),
     ('u', 'user', None, _('list the author (long with -v)')),
@@ -2913,6 +2914,14 @@
 
     Returns 0 if a match is found, 1 otherwise.
     """
+    maxcount = None
+    if opts.get('max_count'):
+        try:
+            maxcount = int(opts.get('max_count'))
+            if maxcount < 1:
+                raise util.Abort(_('invalid max count'))
+        except ValueError:
+            raise util.Abort(_('invalid max count'))
     reflags = re.M
     if opts.get('ignore_case'):
         reflags |= re.I
@@ -2989,6 +2998,13 @@
     def grepbody(fn, rev, body):
         matches[rev].setdefault(fn, [])
         m = matches[rev][fn]
+
+        def addmatch(s):
+            if maxcount and len(m) == maxcount:
+                return False
+            m.append(s)
+            return True
+
         prevlend = 0
         for lnum, lstart, lend, cstart, cend, line in matchlines(body):
             if invert:
@@ -3001,14 +3017,15 @@
                         invertedlstart -= len(inv)
                         s = linestate(inv, invertedlnum, invertedlstart,
                                       invertedlend, -1, -1)
-                        m.append(s)
-
+                        if not addmatch(s):
+                            return
                         invertedlnum += 1
                         invertedlend = invertedlstart - 2
                         invertedlstart = invertedlend
             else:
                 s = linestate(line, lnum, lstart, lend, cstart, cend)
-                m.append(s)
+                if not addmatch(s):
+                    return
             prevlend = lend
         if invert and len(m) and lend != len(body):
             lookahead = body[lend + 1:]
@@ -3020,8 +3037,8 @@
                     invertedlend += len(inv)
                     s = linestate(inv, invertedlnum, invertedlstart,
                                   invertedlend, -1, -1)
-                    m.append(s)
-
+                    if not addmatch(s):
+                        return
                     invertedlnum += 1
                     invertedlstart = invertedlend + 1
                     invertedlend = invertedlstart
diff --git a/tests/test-debugcomplete.t b/tests/test-debugcomplete.t
--- a/tests/test-debugcomplete.t
+++ b/tests/test-debugcomplete.t
@@ -250,7 +250,7 @@
   debugwalk: include, exclude
   debugwireargs: three, four, five, ssh, remotecmd, insecure
   graft: rev, continue, edit, log, currentdate, currentuser, date, user, tool, dry-run
-  grep: print0, all, text, follow, ignore-case, files-with-matches, line-number, no-filename, regexp, invert-match, rev, user, date, include, exclude
+  grep: print0, all, text, follow, ignore-case, files-with-matches, line-number, no-filename, regexp, invert-match, max-count, rev, user, date, include, exclude
   heads: rev, topo, active, closed, style, template
   help: extension, command, keyword
   identify: rev, num, id, branch, tags, bookmarks, ssh, remotecmd, insecure
diff --git a/tests/test-grep.t b/tests/test-grep.t
--- a/tests/test-grep.t
+++ b/tests/test-grep.t
@@ -51,6 +51,21 @@
   $ hg grep -n --invert-match 'ght\nimp' port
   port:4:1:export
 
+maxcount
+  $ hg grep port port -m -1
+  abort: invalid max count
+  [255]
+  $ hg grep port port -m foo
+  abort: invalid max count
+  [255]
+  $ hg grep port port -m 1
+  port:4:export
+  $ hg grep vapor port --invert-match -m 1
+  port:4:export
+  $ hg grep vapor port --invert-match -m 2
+  port:4:export
+  port:4:import/export
+
 simple with color
 
   $ hg --config extensions.color= grep --config color.mode=ansi \