Improve syntax highlighting of multiline strings

Author: xyproto

Diff for: v2/highlight.go

@@ -41,8 +41,6 @@ func (e *Editor) WriteLines(c *vt100.Canvas, fromline, toline LineIndex, cx, cy
 		arrowBeforeCommentMarker           bool
 		inListItem                         bool
 		inCodeBlock                        bool // used when highlighting Doc, Markdown, Python, Nim, Mojo or Starlark
-		threeQuoteStart                    bool
-		threeQuoteEnd                      bool
 		ok                                 bool
 		codeBlockFound                     bool
 		foundDocstringMarker               bool
@@ -155,20 +153,7 @@ func (e *Editor) WriteLines(c *vt100.Canvas, fromline, toline LineIndex, cx, cy
 		for li = LineIndex(0); li < fromline; li++ {
 			line = e.Line(li)
 			trimmedLine = strings.TrimSpace(line)
-			// return """ can be the start of a multiline string, where """ is the end of the multiline string!
-			threeQuoteStart = strings.HasPrefix(trimmedLine, "\"\"\"") || strings.HasPrefix(trimmedLine, "'''")
-			threeQuoteEnd = strings.HasSuffix(trimmedLine, "\"\"\"") || strings.HasSuffix(trimmedLine, "'''")
-			if threeQuoteStart && threeQuoteEnd {
-				inCodeBlock = false
-			} else if trimmedLine == "return \"\"\"" || trimmedLine == "return '''" {
-				inCodeBlock = true
-			} else if strings.HasSuffix(trimmedLine, " = \"\"\"") || strings.HasSuffix(trimmedLine, " = '''") {
-				inCodeBlock = true
-			} else if inCodeBlock && strings.Contains(trimmedLine, "\"\"\"") {
-				inCodeBlock = false
-			} else if threeQuoteStart || threeQuoteEnd {
-				inCodeBlock = !inCodeBlock
-			}
+			inCodeBlock, _ = checkMultiLineString(trimmedLine, inCodeBlock)
 		}
 	}
 

Diff for: v2/strings.go

@@ -228,14 +228,19 @@ func trimRightSpace(str string) string {
 // checkMultiLineString detects and updates the inCodeBlock state.
 // For languages like Nim, Mojo, Python and Starlark.
 func checkMultiLineString(trimmedLine string, inCodeBlock bool) (bool, bool) {
+	trimmedLine = strings.TrimPrefix(trimmedLine, "return ")
 	foundDocstringMarker := false
+	// Check for special syntax patterns that indicate the start of a multiline string
 	if trimmedLine == "\"\"\"" || trimmedLine == "'''" { // only 3 letters
 		inCodeBlock = !inCodeBlock
 		foundDocstringMarker = true
-	} else if strings.HasPrefix(trimmedLine, "\"\"\"") && strings.HasSuffix(trimmedLine, "\"\"\"") { // this could be 6 lte
+	} else if strings.HasSuffix(trimmedLine, " = \"\"\"") || strings.HasSuffix(trimmedLine, " = '''") {
+		inCodeBlock = true
+		foundDocstringMarker = true
+	} else if strings.HasPrefix(trimmedLine, "\"\"\"") && strings.HasSuffix(trimmedLine, "\"\"\"") { // this could be 6 letters
 		inCodeBlock = false
 		foundDocstringMarker = true
-	} else if strings.HasPrefix(trimmedLine, "'''") && strings.HasSuffix(trimmedLine, "'''") { // this could be 6 lettersre
+	} else if strings.HasPrefix(trimmedLine, "'''") && strings.HasSuffix(trimmedLine, "'''") { // this could be 6 letters
 		inCodeBlock = false
 		foundDocstringMarker = true
 	} else if strings.HasPrefix(trimmedLine, "\"\"\"") || strings.HasPrefix(trimmedLine, "'''") { // this is more than 3 ts
@@ -244,7 +249,7 @@ func checkMultiLineString(trimmedLine string, inCodeBlock bool) (bool, bool) {
 			foundDocstringMarker = true
 		}
 	} else if strings.HasSuffix(trimmedLine, "\"\"\"") || strings.HasSuffix(trimmedLine, "'''") { // this is more than 3 ts
-		if (strings.Count(trimmedLine, "\"\"\"") % 2 != 0 || strings.Count(trimmedLine, "'''") % 2 != 0) {
+		if strings.Count(trimmedLine, "\"\"\"")%2 != 0 || strings.Count(trimmedLine, "'''")%2 != 0 {
 			inCodeBlock = !inCodeBlock
 		}
 		if inCodeBlock {