From dad384014b4f2e6a90724e69fa8f26bc99117e9a Mon Sep 17 00:00:00 2001
From: Tulzke <62292054+Tulzke@users.noreply.github.com>
Date: Fri, 15 Sep 2023 19:27:05 +0300
Subject: [PATCH] added ignoring interfaces via -exclude_interfaces flag (#68)
 (#72)

The problem is described in the issue. This is one of the solutions to
this problem.
---
 mockgen/internal/tests/exclude/interfaces.go | 15 ++++++
 mockgen/internal/tests/exclude/mock.go       | 52 +++++++++++++++++++
 mockgen/mockgen.go                           | 19 +++++++
 mockgen/mockgen_test.go                      | 54 ++++++++++++++++++++
 mockgen/parse.go                             |  8 +++
 5 files changed, 148 insertions(+)
 create mode 100644 mockgen/internal/tests/exclude/interfaces.go
 create mode 100644 mockgen/internal/tests/exclude/mock.go

diff --git a/mockgen/internal/tests/exclude/interfaces.go b/mockgen/internal/tests/exclude/interfaces.go
new file mode 100644
index 0000000..216874e
--- /dev/null
+++ b/mockgen/internal/tests/exclude/interfaces.go
@@ -0,0 +1,15 @@
+package exclude
+
+//go:generate mockgen -source=interfaces.go -destination=mock.go -package=ignore -exclude_interfaces=IgnoreMe,IgnoreMe2
+
+type IgnoreMe interface {
+	A() bool
+}
+
+type IgnoreMe2 interface {
+	~int
+}
+
+type GenerateMockForMe interface {
+	B() int
+}
diff --git a/mockgen/internal/tests/exclude/mock.go b/mockgen/internal/tests/exclude/mock.go
new file mode 100644
index 0000000..b8a430f
--- /dev/null
+++ b/mockgen/internal/tests/exclude/mock.go
@@ -0,0 +1,52 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: ./internal/tests/exclude/interfaces.go
+//
+// Generated by this command:
+//
+//	mockgen -source=./internal/tests/exclude/interfaces.go -destination=./internal/tests/exclude/mock.go -package=exclude -exclude_interfaces=IgnoreMe,IgnoreMe2
+//
+// Package exclude is a generated GoMock package.
+package exclude
+
+import (
+	reflect "reflect"
+
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MockGenerateMockForMe is a mock of GenerateMockForMe interface.
+type MockGenerateMockForMe struct {
+	ctrl     *gomock.Controller
+	recorder *MockGenerateMockForMeMockRecorder
+}
+
+// MockGenerateMockForMeMockRecorder is the mock recorder for MockGenerateMockForMe.
+type MockGenerateMockForMeMockRecorder struct {
+	mock *MockGenerateMockForMe
+}
+
+// NewMockGenerateMockForMe creates a new mock instance.
+func NewMockGenerateMockForMe(ctrl *gomock.Controller) *MockGenerateMockForMe {
+	mock := &MockGenerateMockForMe{ctrl: ctrl}
+	mock.recorder = &MockGenerateMockForMeMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockGenerateMockForMe) EXPECT() *MockGenerateMockForMeMockRecorder {
+	return m.recorder
+}
+
+// B mocks base method.
+func (m *MockGenerateMockForMe) B() int {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "B")
+	ret0, _ := ret[0].(int)
+	return ret0
+}
+
+// B indicates an expected call of B.
+func (mr *MockGenerateMockForMeMockRecorder) B() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "B", reflect.TypeOf((*MockGenerateMockForMe)(nil).B))
+}
diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go
index a359c3c..feb747f 100644
--- a/mockgen/mockgen.go
+++ b/mockgen/mockgen.go
@@ -65,6 +65,7 @@ var (
 	typed                  = flag.Bool("typed", false, "Generate Type-safe 'Return', 'Do', 'DoAndReturn' function")
 	imports                = flag.String("imports", "", "(source mode) Comma-separated name=path pairs of explicit imports to use.")
 	auxFiles               = flag.String("aux_files", "", "(source mode) Comma-separated pkg=path pairs of auxiliary Go source files.")
+	excludeInterfaces      = flag.String("exclude_interfaces", "", "Comma-separated names of interfaces to be excluded")
 
 	debugParser = flag.Bool("debug_parser", false, "Print out parser results only.")
 	showVersion = flag.Bool("version", false, "Print version.")
@@ -200,6 +201,24 @@ func parseMockNames(names string) map[string]string {
 	return mocksMap
 }
 
+func parseExcludeInterfaces(names string) map[string]struct{} {
+	splitNames := strings.Split(names, ",")
+	namesSet := make(map[string]struct{}, len(splitNames))
+	for _, name := range splitNames {
+		if name == "" {
+			continue
+		}
+
+		namesSet[name] = struct{}{}
+	}
+
+	if len(namesSet) == 0 {
+		return nil
+	}
+
+	return namesSet
+}
+
 func usage() {
 	_, _ = io.WriteString(os.Stderr, usageText)
 	flag.PrintDefaults()
diff --git a/mockgen/mockgen_test.go b/mockgen/mockgen_test.go
index 3bd9fc3..6ac70dd 100644
--- a/mockgen/mockgen_test.go
+++ b/mockgen/mockgen_test.go
@@ -412,3 +412,57 @@ func TestParsePackageImport_FallbackMultiGoPath(t *testing.T) {
 		t.Errorf("expect %s, got %s", expectedPkgPath, pkgPath)
 	}
 }
+
+func TestParseExcludeInterfaces(t *testing.T) {
+	testCases := []struct {
+		name     string
+		arg      string
+		expected map[string]struct{}
+	}{
+		{
+			name:     "empty string",
+			arg:      "",
+			expected: nil,
+		},
+		{
+			name:     "string without a comma",
+			arg:      "arg1",
+			expected: map[string]struct{}{"arg1": {}},
+		},
+		{
+			name:     "two names",
+			arg:      "arg1,arg2",
+			expected: map[string]struct{}{"arg1": {}, "arg2": {}},
+		},
+		{
+			name:     "two names with a comma at the end",
+			arg:      "arg1,arg2,",
+			expected: map[string]struct{}{"arg1": {}, "arg2": {}},
+		},
+		{
+			name:     "two names with a comma at the beginning",
+			arg:      ",arg1,arg2",
+			expected: map[string]struct{}{"arg1": {}, "arg2": {}},
+		},
+		{
+			name:     "commas only",
+			arg:      ",,,,",
+			expected: nil,
+		},
+		{
+			name:     "duplicates",
+			arg:      "arg1,arg2,arg1",
+			expected: map[string]struct{}{"arg1": {}, "arg2": {}},
+		},
+	}
+
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			actual := parseExcludeInterfaces(tt.arg)
+
+			if !reflect.DeepEqual(actual, tt.expected) {
+				t.Errorf("expected %v, actual %v", tt.expected, actual)
+			}
+		})
+	}
+}
diff --git a/mockgen/parse.go b/mockgen/parse.go
index 8379fa3..9521409 100644
--- a/mockgen/parse.go
+++ b/mockgen/parse.go
@@ -75,6 +75,10 @@ func sourceMode(source string) (*model.Package, error) {
 		}
 	}
 
+	if *excludeInterfaces != "" {
+		p.excludeNamesSet = parseExcludeInterfaces(*excludeInterfaces)
+	}
+
 	// Handle -aux_files.
 	if err := p.parseAuxFiles(*auxFiles); err != nil {
 		return nil, err
@@ -163,6 +167,7 @@ type fileParser struct {
 	auxFiles           []*ast.File
 	auxInterfaces      *interfaceCache
 	srcDir             string
+	excludeNamesSet    map[string]struct{}
 }
 
 func (p *fileParser) errorf(pos token.Pos, format string, args ...any) error {
@@ -223,6 +228,9 @@ func (p *fileParser) parseFile(importPath string, file *ast.File) (*model.Packag
 
 	var is []*model.Interface
 	for ni := range iterInterfaces(file) {
+		if _, ok := p.excludeNamesSet[ni.name.String()]; ok {
+			continue
+		}
 		i, err := p.parseInterface(ni.name.String(), importPath, ni)
 		if err != nil {
 			return nil, err