Repository: bombsimon/wsl
Pull request #103 (2021-04-19 21:04:00)

#102 allow forcing short declarations to cuddle only with themselves
This option treats a short declaration with := as different than a plain
assignment with = and so prevents them cuddling together, as well as not
allowing a short declaration to cuddle with a block, for example.

mholiday-nyt: > Thank you very much for this, nice implemented and appreciated the really through documentation! I added a question about a node (if conversion is needed) and a suggestion to simplify a method. Let me know what you think and I'll be ready to merge!

With these code review changes, I think we're good

Patch:
diff --git a/README.md b/README.md
index c5068a8..9812f94 100644
--- a/README.md
+++ b/README.md
@@ -121,5 +121,6 @@ feel free to raise an [issue](https://github.com/bombsimon/wsl/issues/new).
 * [Only one cuddle assignment allowed before type switch statement](doc/rules.md#only-one-cuddle-assignment-allowed-before-type-switch-statement)
 * [Ranges should only be cuddled with assignments used in the iteration](doc/rules.md#ranges-should-only-be-cuddled-with-assignments-used-in-the-iteration)
 * [Return statements should not be cuddled if block has more than two lines](doc/rules.md#return-statements-should-not-be-cuddled-if-block-has-more-than-two-lines)
+* [Short declarations should cuddle only with other short declarations](doc/rules.md#short-declaration-should-cuddle-only-with-other-short-declarations)
 * [Switch statements should only be cuddled with variables switched](doc/rules.md#switch-statements-should-only-be-cuddled-with-variables-switched)
 * [Type switch statements should only be cuddled with variables switched](doc/rules.md#type-switch-statements-should-only-be-cuddled-with-variables-switched)
diff --git a/cmd/wsl/main.go b/cmd/wsl/main.go
index 85dc070..414cb7c 100644
--- a/cmd/wsl/main.go
+++ b/cmd/wsl/main.go
@@ -38,6 +38,7 @@ func main() {
 	flag.BoolVar(&config.AllowTrailingComment, "allow-trailing-comment", false, "Allow blocks to end with a comment")
 	flag.BoolVar(&config.AllowSeparatedLeadingComment, "allow-separated-leading-comment", false, "Allow empty newlines in leading comments")
 	flag.BoolVar(&config.ForceCuddleErrCheckAndAssign, "force-err-cuddling", false, "Force cuddling of error checks with error var assignment")
+	flag.BoolVar(&config.ForceExclusiveShortDeclarations, "force-short-decl-cuddling", false, "Force short declarations to cuddle by themselves")
 	flag.IntVar(&config.ForceCaseTrailingWhitespaceLimit, "force-case-trailing-whitespace", 0, "Force newlines for case blocks > this number.")
 
 	flag.Parse()
diff --git a/doc/configuration.md b/doc/configuration.md
index 689524e..f2ec62e 100644
--- a/doc/configuration.md
+++ b/doc/configuration.md
@@ -296,3 +296,41 @@ if err != nil {
     return err
 }
 ```
+
+### [force-short-decl-cuddling](rules.md#short-declaration-should-cuddle-only-with-other-short-declarations)
+
+Enforces that an assignment which is actually a short declaration (using `:=`)
+is only allowed to cuddle with other short declarations, and not plain
+assignments, blocks, etc. This rule helps make declarations stand out by
+themselves, much the same as grouping `var` statements.
+
+> Default value: false
+
+Supported when false:
+
+```go
+a := 1
+a = 2
+
+err := ErrorProducingFunc()
+if err != nil {
+    return err
+}
+```
+
+Required when true:
+
+```go
+a := 1
+
+a = 2
+
+err := ErrorProducingFunc()
+
+if err != nil {
+    return err
+}
+```
+
+**Note**: this means the option _overrides_ the 
+[enforce-err-cuddling](#enforce-err-cuddling) option above, among others.
diff --git a/doc/rules.md b/doc/rules.md
index 1c81d82..d80afff 100644
--- a/doc/rules.md
+++ b/doc/rules.md
@@ -1013,3 +1013,38 @@ if err != nil {
     return err
 }
 ```
+
+---
+
+### Short declaration should cuddle only with other short declarations
+
+> Can be configured, see [configuration
+documentation](configuration.md#force-short-decl-cuddling)
+
+A short declaration (an "assignment" using `:=`) must not be cuddled with anything other than another short declaration.
+
+```go
+a := 1
+err := ErrorProducingFunc()
+if err != nil {
+    return err
+}
+```
+
+#### Recommended amendment
+
+Add an empty line between the short declaration and the error check.
+
+```go
+a := 1
+err := ErrorProducingFunc()
+
+if err != nil {
+    return err
+}
+```
+
+**Note**: this is the opposite of the case above forcing an `err` variable
+to be assigned together with its check; it also overrides some other rules
+which about assignments by separating short declarations from "plain"
+assignment statements using `=`.
diff --git a/wsl.go b/wsl.go
index ae0dfb4..9389f4e 100644
--- a/wsl.go
+++ b/wsl.go
@@ -41,6 +41,7 @@ const (
 	reasonBlockStartsWithWS              = "block should not start with a whitespace"
 	reasonBlockEndsWithWS                = "block should not end with a whitespace (or comment)"
 	reasonCaseBlockTooCuddly             = "case block should end with newline at this size"
+	reasonShortDeclNotExclusive          = "short declaration should cuddle only with other short declarations"
 )
 
 // Warning strings
@@ -159,6 +160,20 @@ type Configuration struct {
 	// used for error variables to check for in the conditional.
 	// Defaults to just "err"
 	ErrorVariableNames []string
+
+	// ForceExclusiveShortDeclarations will cause an error if a short declaration
+	// (:=) cuddles with anything other than another short declaration. For example
+	//
+	// a := 2
+	// b := 3
+	//
+	// is allowed, but
+	//
+	// a := 2
+	// b = 3
+	//
+	// is not allowed. This logic overrides ForceCuddleErrCheckAndAssign among others.
+	ForceExclusiveShortDeclarations bool
 }
 
 // DefaultConfig returns default configuration
@@ -171,6 +186,7 @@ func DefaultConfig() Configuration {
 		AllowTrailingComment:             false,
 		AllowSeparatedLeadingComment:     false,
 		ForceCuddleErrCheckAndAssign:     false,
+		ForceExclusiveShortDeclarations:  false,
 		ForceCaseTrailingWhitespaceLimit: 0,
 		AllowCuddleWithCalls:             []string{"Lock", "RLock"},
 		AllowCuddleWithRHS:               []string{"Unlock", "RUnlock"},
@@ -391,6 +407,18 @@ func (p *Processor) parseBlockStatements(statements []ast.Stmt) {
 			return false
 		}
 
+		// If it's a short declaration we should not cuddle with anything else
+		// if ForceExclusiveShortDeclarations is set on; either this or the
+		// previous statement could be the short decl, so we'll find out which
+		// it was and use *that* statement's position
+		if p.config.ForceExclusiveShortDeclarations && cuddledWithLastStmt {
+			if p.isShortDecl(stmt) && !p.isShortDecl(previousStatement) {
+				p.addError(stmt.Pos(), reasonShortDeclNotExclusive)
+			} else if p.isShortDecl(previousStatement) && !p.isShortDecl(stmt) {
+				p.addError(previousStatement.Pos(), reasonShortDeclNotExclusive)
+			}
+		}
+
 		// If it's not an if statement and we're not cuddled move on. The only
 		// reason we need to keep going for if statements is to check if we
 		// should be cuddled with an error check.
@@ -888,6 +916,14 @@ func (p *Processor) findRHS(node ast.Node) []string {
 	return rhs
 }
 
+func (p *Processor) isShortDecl(node ast.Node) bool {
+	if t, ok := node.(*ast.AssignStmt); ok {
+		return t.Tok == token.DEFINE
+	}
+
+	return false
+}
+
 func (p *Processor) findBlockStmt(node ast.Node) []*ast.BlockStmt {
 	var blocks []*ast.BlockStmt
 


Test patch:
diff --git a/wsl_test.go b/wsl_test.go
index ffa361c..41d62d5 100644
--- a/wsl_test.go
+++ b/wsl_test.go
@@ -1828,6 +1828,35 @@ func TestWithConfig(t *testing.T) {
 				reasonOnlyOneCuddle,
 			},
 		},
+		{
+			description: "force short decls to cuddle alone",
+			code: []byte(`package main
+
+			func main() {
+				// OK
+				a := 1
+				b := 2
+
+				// should fail
+				c := 3
+				b = 2
+
+				// should fail 
+				err := DoSomething()
+				if err != nil {
+					b = 4
+				}
+			}`),
+			customConfig: &Configuration{
+				AllowAssignAndAnythingCuddle:    true,
+				ForceCuddleErrCheckAndAssign:    true,
+				ForceExclusiveShortDeclarations: true,
+			},
+			expectedErrorStrings: []string{
+				reasonShortDeclNotExclusive,
+				reasonShortDeclNotExclusive,
+			},
+		},
 	}
 
 	for _, tc := range cases {