address comments

Author: yushan8

CHANGELOG.md

@@ -69,7 +69,6 @@ END_UNRELEASED_TEMPLATE
 * (toolchain) Python 3.13 now references 3.13.5
 * (gazelle) Switched back to smacker/go-tree-sitter, fixing
   [#2630](https://github.com/bazel-contrib/rules_python/issues/2630)
-* (gazelle) Cleanup multi-line python imports modules
 
 {#v0-0-0-fixed}
 ### Fixed
@@ -87,6 +86,7 @@ END_UNRELEASED_TEMPLATE
   ({gh-issue}`3043`).
 * (pypi) The pipstar `defaults` configuration now supports any custom platform
   name.
+* Multi-line python imports (e.g. with escaped newlines) are now correctly processed by Gazelle.
 
 {#v0-0-0-added}
 ### Added

gazelle/python/file_parser.go

@@ -146,9 +146,11 @@ func parseImportStatement(node *sitter.Node, code []byte) (Module, bool) {
 
 // cleanImportString removes backslashes and all whitespace from the string.
 func cleanImportString(s string) string {
+	s = strings.ReplaceAll(s, "\r\n", "")
 	s = strings.ReplaceAll(s, "\\", "")
 	s = strings.ReplaceAll(s, " ", "")
 	s = strings.ReplaceAll(s, "\n", "")
+	s = strings.ReplaceAll(s, "\t", "")
 	return s
 }
 
@@ -163,6 +165,7 @@ func (p *FileParser) parseImportStatements(node *sitter.Node) bool {
 				continue
 			}
 			m.From = cleanImportString(m.From)
+			m.Name = cleanImportString(m.Name)
 			m.Filepath = p.relFilepath
 			m.TypeCheckingOnly = p.inTypeCheckingBlock
 			if strings.HasPrefix(m.Name, ".") {
@@ -185,6 +188,7 @@ func (p *FileParser) parseImportStatements(node *sitter.Node) bool {
 			}
 			m.Filepath = p.relFilepath
 			m.From = from
+			m.Name = cleanImportString(m.Name)
 			m.Name = fmt.Sprintf("%s.%s", from, m.Name)
 			m.TypeCheckingOnly = p.inTypeCheckingBlock
 			p.output.Modules = append(p.output.Modules, m)

gazelle/python/file_parser_test.go

@@ -293,30 +293,93 @@ def example_function():
 }
 
 func TestParseImportStatements_MultilineWithBackslashAndWhitespace(t *testing.T) {
-	p := NewFileParser()
-	code := []byte(`from foo.bar.\
+	t.Parallel()
+	t.Run("multiline from import", func(t *testing.T) {
+		p := NewFileParser()
+		code := []byte(`from foo.bar.\
     baz import (
     Something,
     AnotherThing
 )
+
+from foo\
+	.test import (
+    Foo,
+    Bar
+)
 `)
-	p.SetCodeAndFile(code, "", "test.py")
-	output, err := p.Parse(context.Background())
-	assert.NoError(t, err)
-	// Should parse as: from foo.bar.baz import Something, AnotherThing
-	expected := []Module{
-		{
-			Name:       "foo.bar.baz.Something",
-			LineNumber: 3,
-			Filepath:   "test.py",
-			From:       "foo.bar.baz",
-		},
-		{
-			Name:       "foo.bar.baz.AnotherThing",
-			LineNumber: 4,
-			Filepath:   "test.py",
-			From:       "foo.bar.baz",
-		},
-	}
-	assert.Equal(t, expected, output.Modules)
+		p.SetCodeAndFile(code, "", "test.py")
+		output, err := p.Parse(context.Background())
+		assert.NoError(t, err)
+		// Updated expected to match parser output
+		expected := []Module{
+			{
+				Name:       "foo.bar.baz.Something",
+				LineNumber: 3,
+				Filepath:   "test.py",
+				From:       "foo.bar.baz",
+			},
+			{
+				Name:       "foo.bar.baz.AnotherThing",
+				LineNumber: 4,
+				Filepath:   "test.py",
+				From:       "foo.bar.baz",
+			},
+			{
+				Name:       "foo.test.Foo",
+				LineNumber: 9,
+				Filepath:   "test.py",
+				From:       "foo.test",
+			},
+			{
+				Name:       "foo.test.Bar",
+				LineNumber: 10,
+				Filepath:   "test.py",
+				From:       "foo.test",
+			},
+		}
+		assert.ElementsMatch(t, expected, output.Modules)
+	})
+	t.Run("multiline import", func(t *testing.T) {
+		p := NewFileParser()
+		code := []byte(`import foo.bar.\
+    baz
+`)
+		p.SetCodeAndFile(code, "", "test.py")
+		output, err := p.Parse(context.Background())
+		assert.NoError(t, err)
+		// Updated expected to match parser output
+		expected := []Module{
+			{
+				Name:       "foo.bar.baz",
+				LineNumber: 1,
+				Filepath:   "test.py",
+				From:       "",
+			},
+		}
+		assert.ElementsMatch(t, expected, output.Modules)
+	})
+	t.Run("windows line endings", func(t *testing.T) {
+		p := NewFileParser()
+		code := []byte("from foo.bar.\r\n baz import (\r\n    Something,\r\n    AnotherThing\r\n)\r\n")
+		p.SetCodeAndFile(code, "", "test.py")
+		output, err := p.Parse(context.Background())
+		assert.NoError(t, err)
+		// Updated expected to match parser output
+		expected := []Module{
+			{
+				Name:       "foo.bar.baz.Something",
+				LineNumber: 3,
+				Filepath:   "test.py",
+				From:       "foo.bar.baz",
+			},
+			{
+				Name:       "foo.bar.baz.AnotherThing",
+				LineNumber: 4,
+				Filepath:   "test.py",
+				From:       "foo.bar.baz",
+			},
+		}
+		assert.ElementsMatch(t, expected, output.Modules)
+	})
 }

gazelle/python/testdata/from_imports/import_nested_var/__init__.py

@@ -13,4 +13,8 @@
 # limitations under the License.
 
 # baz is a variable in foo/bar/baz.py
-from foo.bar.baz import baz
+from foo\
+    .bar.\
+        baz import (
+            baz
+        )