Equip Config with a Stoplist

hrs · May 23, 2023 · 6d73861 · 6d73861
1 parent aedb7c2
commit 6d73861
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 6 deletions.
diff --git a/lib/config.go b/lib/config.go
@@ -8,5 +8,6 @@ type Config struct {
 	NoStoplist     bool
 	OmitQuery      bool
 	ShowScores     bool
+	Stoplist       *Stoplist
 	Verbose        bool
 }
diff --git a/lib/document.go b/lib/document.go
@@ -60,7 +60,7 @@ func NewDocument(path string, config *Config) (*Document, error) {
 			word = strings.TrimSuffix(word, "'s")
 
 			if word != "" {
-				if config.NoStoplist || !InStoplist(word) {
+				if config.NoStoplist || !config.Stoplist.Include(word) {
 					if config.NoStemming {
 						termCount[word]++
 					} else {

diff --git a/lib/main.go b/lib/main.go
@@ -28,6 +28,7 @@ func main() {
 		NoStoplist:     *noStoplistFlag,
 		OmitQuery:      *omitQueryFlag,
 		ShowScores:     *showScoresFlag,
+		Stoplist:       &DefaultStoplist,
 		Verbose:        *verboseFlag,
 	}
 

diff --git a/lib/stoplist.go b/lib/stoplist.go
@@ -1,11 +1,23 @@
 package main
 
-func InStoplist(word string) bool {
-	_, ok := stoplist[word]
+type Stoplist map[string]bool
+
+func (stoplist *Stoplist) Include(term string) bool {
+	_, ok := (*stoplist)[term]
 	return ok
 }
 
-var stoplist = map[string]bool{
+func NewStoplist(terms []string) *Stoplist {
+	stoplist := make(Stoplist)
+
+	for _, term := range terms {
+		stoplist[term] = true
+	}
+
+	return &stoplist
+}
+
+var DefaultStoplist = Stoplist{
 	"'ll":             true,
 	"'ve":             true,
 	"I":               true,

diff --git a/lib/stoplist_test.go b/lib/stoplist_test.go
@@ -2,7 +2,7 @@ package main
 
 import "testing"
 
-func TestInStoplist(t *testing.T) {
+func TestDefaultStoplist(t *testing.T) {
 	tests := []struct {
 		word     string
 		expected bool
@@ -14,7 +14,31 @@ func TestInStoplist(t *testing.T) {
 	}
 
 	for _, tc := range tests {
-		got := InStoplist(tc.word)
+		got := DefaultStoplist.Include(tc.word)
+
+		if got != tc.expected {
+			t.Errorf("got %t, wanted %t", got, tc.expected)
+		}
+	}
+}
+
+func TestCustomStoplist(t *testing.T) {
+	stoplist := NewStoplist([]string{
+		"foo",
+		"bar",
+	})
+
+	tests := []struct {
+		word     string
+		expected bool
+	}{
+		{"foo", true},
+		{"bar", true},
+		{"baz", false},
+	}
+
+	for _, tc := range tests {
+		got := stoplist.Include(tc.word)
 
 		if got != tc.expected {
 			t.Errorf("got %t, wanted %t", got, tc.expected)