Allow multiline assignments to be cuddled. Support config flags

cmd/wsl/main.go

@@ -19,6 +19,7 @@ func main() {
 		cwd, _       = os.Getwd()
 		files        = []string{}
 		finalFiles   = []string{}
+		config       = wsl.DefaultConfig()
 	)
 
 	flag.BoolVar(&help, "h", false, "Show this help text")
@@ -28,6 +29,10 @@ func main() {
 	flag.BoolVar(&showWarnings, "w", false, "Show warnings (ignored rules)")
 	flag.BoolVar(&showWarnings, "warnings", false, "")
 
+	flag.BoolVar(&config.CheckAppend, "check-append", true, "Strict rules for append")
+	flag.BoolVar(&config.AllowAssignAndCallsCuddle, "allow-assign-and-call", true, "Allow assignments and calls to be cuddled (if using same variable/type)")
+	flag.BoolVar(&config.AllowMultiLineAssignmentCuddled, "allow-multi-line-assignment", true, "Allow cuddling with multi line assignments")
+
 	flag.Parse()
 
 	if help {
@@ -71,7 +76,7 @@ func main() {
 		finalFiles = append(finalFiles, f)
 	}
 
-	processor := wsl.NewProcessor()
+	processor := wsl.NewProcessorWithConfig(config)
 	result, warnings := processor.ProcessFiles(finalFiles)
 
 	for _, r := range result {

wsl.go

@@ -26,6 +26,17 @@ type Configuration struct {
 	//  y = AnotherAssign()
 	AllowAssignAndCallsCuddle bool
 
+	// AllowMultiLineAssignmentCuddled allows cuddling to assignments even if
+	// they span over multiple lines. This defaults to true which allows the
+	// following example:
+	//  err := function(
+	//  	"multiple", "lines",
+	//  )
+	//  if err != nil {
+	//  	// ...
+	//  }
+	AllowMultiLineAssignmentCuddled bool
+
 	// AllowCuddleWithCalls is a list of call idents that everything can be
 	// cuddled with. Defaults to calls looking like locks to support a flow like
 	// this:
@@ -41,6 +52,17 @@ type Configuration struct {
 	AllowCuddleWithRHS []string
 }
 
+// DefaultConfig returns default configuration
+func DefaultConfig() Configuration {
+	return Configuration{
+		CheckAppend:                     true,
+		AllowAssignAndCallsCuddle:       true,
+		AllowMultiLineAssignmentCuddled: true,
+		AllowCuddleWithCalls:            []string{"Lock", "RLock"},
+		AllowCuddleWithRHS:              []string{"Unlock", "RUnlock"},
+	}
+}
+
 type Result struct {
 	FileName   string
 	LineNumber int
@@ -70,12 +92,7 @@ func NewProcessorWithConfig(cfg Configuration) *Processor {
 
 // NewProcessor will create a Processor.
 func NewProcessor() *Processor {
-	return NewProcessorWithConfig(Configuration{
-		CheckAppend:               true,
-		AllowAssignAndCallsCuddle: true,
-		AllowCuddleWithCalls:      []string{"Lock", "RLock", "RWLock"},
-		AllowCuddleWithRHS:        []string{"Unlock", "RUnlock", "RWUnlock"},
-	})
+	return NewProcessorWithConfig(DefaultConfig())
 }
 
 // ProcessFiles takes a string slice with file names (full paths) and lints
@@ -141,6 +158,7 @@ func (p *Processor) parseBlockBody(block *ast.BlockStmt) {
 
 // parseBlockStatements will parse all the statements found in the body of a
 // node. A list of Result is returned.
+// nolint: gocognit
 func (p *Processor) parseBlockStatements(statements []ast.Stmt) {
 	for i, stmt := range statements {
 		// TODO: How to tell when and where func literals may exist to enforce
@@ -177,13 +195,24 @@ func (p *Processor) parseBlockStatements(statements []ast.Stmt) {
 		// special handling of things such as mutexes.
 		var calledOnLineAbove []string
 
+		// Check if the previous statement spans over multiple lines.
+		var isMultiLineAssignment = p.nodeStart(previousStatement) != p.nodeStart(stmt)-1
+
 		// Ensure previous line is not a multi line assignment and if not get
 		// rightAndLeftHandSide assigned variables.
-		if p.nodeStart(previousStatement) == p.nodeStart(stmt)-1 {
+		if !isMultiLineAssignment {
 			assignedOnLineAbove = p.findLHS(previousStatement)
 			calledOnLineAbove = p.findRHS(previousStatement)
 		}
 
+		// If previous assignment is multi line and we allow it, fetch
+		// assignments (but only assignments).
+		if isMultiLineAssignment && p.config.AllowMultiLineAssignmentCuddled {
+			if _, ok := previousStatement.(*ast.AssignStmt); ok {
+				assignedOnLineAbove = p.findLHS(previousStatement)
+			}
+		}
+
 		// We could potentially have a block which require us to check the first
 		// argument before ruling out an allowed cuddle.
 		var assignedFirstInBlock []string

wsl_test.go

@@ -407,23 +407,6 @@ func TestShouldAddEmptyLines(t *testing.T) {
 				}
 			}`),
 		},
-		{
-			description: "multi line assignment should not be seen as assignment above",
-			code: []byte(`package main
-
-			func main() {
-				// This should not be OK since the assignment is split at
-				// multiple rows.
-				err := SomeFunc(
-					"taking multiple values",
-					"suddendly not a single line",
-				)
-				if err != nil {
-					fmt.Println("denied")
-				}
-			}`),
-			expectedErrorStrings: []string{"if statements should only be cuddled with assignments"},
-		},
 		{
 			description: "allow cuddle for immediate assignment in block",
 			code: []byte(`package main
@@ -1021,6 +1004,55 @@ func TestWithConfig(t *testing.T) {
 				AllowAssignAndCallsCuddle: true,
 			},
 		},
+		{
+			description: "multi line assignment ok when enabled",
+			code: []byte(`package main
+
+			func main() {
+				// This should not be OK since the assignment is split at
+				// multiple rows.
+				err := SomeFunc(
+					"taking multiple values",
+					"suddendly not a single line",
+				)
+				if err != nil {
+					fmt.Println("denied")
+				}
+			}`),
+			customConfig: &Configuration{
+				AllowMultiLineAssignmentCuddled: false,
+			},
+			expectedErrorStrings: []string{"if statements should only be cuddled with assignments"},
+		},
+		{
+			description: "multi line assignment not ok when disabled",
+			code: []byte(`package main
+
+			func main() {
+				// This should not be OK since the assignment is split at
+				// multiple rows.
+				err := SomeFunc(
+					"taking multiple values",
+					"suddenly not a single line",
+				)
+				if err != nil {
+					fmt.Println("accepted")
+				}
+
+				// This should still fail since it's not the same variable
+				noErr := SomeFunc(
+					"taking multiple values",
+					"suddenly not a single line",
+				)
+				if err != nil {
+					fmt.Println("rejected")
+				}
+			}`),
+			customConfig: &Configuration{
+				AllowMultiLineAssignmentCuddled: true,
+			},
+			expectedErrorStrings: []string{"if statements should only be cuddled with assignments used in the if statement itself"},
+		},
 	}
 
 	for _, tc := range cases {