github.com/anonymouse64/snapd@v0.0.0-20210824153203-04c4c42d842d/strutil/matchcounter.go (about) 1 // -*- Mode: Go; indent-tabs-mode: t -*- 2 3 /* 4 * Copyright (C) 2018 Canonical Ltd 5 * 6 * This program is free software: you can redistribute it and/or modify 7 * it under the terms of the GNU General Public License version 3 as 8 * published by the Free Software Foundation. 9 * 10 * This program is distributed in the hope that it will be useful, 11 * but WITHOUT ANY WARRANTY; without even the implied warranty of 12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 13 * GNU General Public License for more details. 14 * 15 * You should have received a copy of the GNU General Public License 16 * along with this program. If not, see <http://www.gnu.org/licenses/>. 17 * 18 */ 19 20 package strutil 21 22 import ( 23 "bytes" 24 "regexp" 25 ) 26 27 // A MatchCounter is a discarding io.Writer that retains up to N 28 // matches to its Regexp before just counting matches. 29 // 30 // It does not work with regexps that cross newlines; in fact it will 31 // probably not work if the data written isn't line-orineted. 32 // 33 // If Regexp is not set (or nil), it matches whole non-empty lines. 34 type MatchCounter struct { 35 // Regexp to use to find matches in the stream 36 Regexp *regexp.Regexp 37 // Maximum number of matches to keep; if < 0, keep all matches 38 N int 39 // LastN when true indicates that only N last matches shall be kept. 40 LastN bool 41 42 count int 43 matches []string 44 partial bytes.Buffer 45 } 46 47 func (w *MatchCounter) Write(p []byte) (int, error) { 48 n := len(p) 49 if w.partial.Len() > 0 { 50 idx := bytes.IndexByte(p, '\n') 51 if idx < 0 { 52 // no newline yet, carry on accumulating 53 w.partial.Write(p) 54 return n, nil 55 } 56 idx++ 57 w.partial.Write(p[:idx]) 58 w.check(w.partial.Bytes()) 59 p = p[idx:] 60 } 61 w.partial.Reset() 62 idx := bytes.LastIndexByte(p, '\n') 63 if idx < 0 { 64 w.partial.Write(p) 65 return n, nil 66 } 67 idx++ 68 w.partial.Write(p[idx:]) 69 w.check(p[:idx]) 70 return n, nil 71 } 72 73 func (w *MatchCounter) check(p []byte) { 74 addMatch := func(m string) { 75 switch { 76 case w.N == 0: 77 // keep none 78 return 79 case w.N < 0: 80 // keep all 81 fallthrough 82 case !w.LastN && len(w.matches) < w.N: 83 w.matches = append(w.matches, m) 84 case w.LastN: 85 // keep only last N matches 86 keep := w.matches 87 if len(w.matches)+1 > w.N { 88 keep = w.matches[1:] 89 } 90 w.matches = append(keep, m) 91 } 92 } 93 if w.Regexp == nil { 94 for { 95 idx := bytes.IndexByte(p, '\n') 96 if idx < 0 { 97 return 98 } 99 if idx == 0 { 100 // empty line 101 p = p[1:] 102 continue 103 } 104 addMatch(string(p[:idx])) 105 w.count++ 106 p = p[idx+1:] 107 } 108 } 109 matches := w.Regexp.FindAll(p, -1) 110 for _, match := range matches { 111 addMatch(string(match)) 112 } 113 w.count += len(matches) 114 } 115 116 // Matches returns the first few matches, and the total number of matches seen. 117 func (w *MatchCounter) Matches() ([]string, int) { 118 return w.matches, w.count 119 }