Address PR #19 review: defensive checks and TYPE_CHECKING tests

- parseMain: add ChildCount and nil checks before indexing into
  comparison operator children to prevent panics on malformed ASTs
- parseImportStatements: guard import_from_statement with ChildCount
  check to handle incomplete/invalid syntax safely
- Add TestParseTypeChecking covering TYPE_CHECKING, typing.TYPE_CHECKING,
  and non-TYPE_CHECKING conditional import blocks

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ewianda

python/file_parser.go

@@ -105,20 +105,27 @@ func parseCode(code []byte) (*sitter.Node, error) {
 func (p *FileParser) parseMain(node *sitter.Node) bool {
 	for i := 0; i < int(node.ChildCount()); i++ {
 		child := node.Child(i)
-		if child.Type() == "if_statement" &&
-			child.ChildCount() > 1 &&
-			child.Child(1).Type() == sitterNodeTypeComparisonOperator &&
-			child.Child(1).Child(1).Type() == "==" {
-			statement := child.Child(1)
-			a, b := statement.Child(0), statement.Child(2)
-			// convert "'__main__' == __name__" to "__name__ == '__main__'"
-			if b.Type() == sitterNodeTypeIdentifier {
-				a, b = b, a
-			}
-			if a.Type() == sitterNodeTypeIdentifier && a.Content(p.code) == "__name__" &&
-				b.Type() == sitterNodeTypeString && string(p.code[b.StartByte()+1:b.EndByte()-1]) == "__main__" {
-				return true
-			}
+		if child.Type() != "if_statement" || child.ChildCount() < 2 {
+			continue
+		}
+		cond := child.Child(1)
+		if cond == nil || cond.Type() != sitterNodeTypeComparisonOperator || cond.ChildCount() < 3 {
+			continue
+		}
+		if cond.Child(1).Type() != "==" {
+			continue
+		}
+		a, b := cond.Child(0), cond.Child(2)
+		if a == nil || b == nil {
+			continue
+		}
+		// convert "'__main__' == __name__" to "__name__ == '__main__'"
+		if b.Type() == sitterNodeTypeIdentifier {
+			a, b = b, a
+		}
+		if a.Type() == sitterNodeTypeIdentifier && a.Content(p.code) == "__name__" &&
+			b.Type() == sitterNodeTypeString && string(p.code[b.StartByte()+1:b.EndByte()-1]) == "__main__" {
+			return true
 		}
 	}
 	return false
@@ -174,6 +181,9 @@ func (p *FileParser) parseImportStatements(node *sitter.Node) {
 			p.output.Modules = append(p.output.Modules, m)
 		}
 	} else if node.Type() == sitterNodeTypeImportFromStatement {
+		if node.ChildCount() < 4 {
+			return
+		}
 		from := node.Child(1).Content(p.code)
 		if strings.HasPrefix(from, ".") {
 			return

python/file_parser_test.go

@@ -254,3 +254,60 @@ func TestParseFull(t *testing.T) {
 		FileName: "a.py",
 	}, *output)
 }
+
+func TestParseTypeChecking(t *testing.T) {
+	t.Parallel()
+	units := []struct {
+		name     string
+		code     string
+		filepath string
+		result   []module
+	}{
+		{
+			name: "TYPE_CHECKING bare",
+			code: `import os
+if TYPE_CHECKING:
+    from foo import bar
+`,
+			filepath: "abc.py",
+			result: []module{
+				{Name: "os", LineNumber: 1, Filepath: "abc.py"},
+				{Name: "foo.bar", LineNumber: 3, Filepath: "abc.py", From: "foo"},
+			},
+		},
+		{
+			name: "typing.TYPE_CHECKING",
+			code: `import os
+if typing.TYPE_CHECKING:
+    from baz import qux
+    import something
+`,
+			filepath: "abc.py",
+			result: []module{
+				{Name: "os", LineNumber: 1, Filepath: "abc.py"},
+				{Name: "baz.qux", LineNumber: 3, Filepath: "abc.py", From: "baz"},
+				{Name: "something", LineNumber: 4, Filepath: "abc.py"},
+			},
+		},
+		{
+			name: "not TYPE_CHECKING",
+			code: `if some_other_condition:
+    from foo import bar
+`,
+			filepath: "abc.py",
+			result: []module{
+				{Name: "foo.bar", LineNumber: 2, Filepath: "abc.py", From: "foo"},
+			},
+		},
+	}
+	for _, u := range units {
+		t.Run(u.name, func(t *testing.T) {
+			p := NewFileParser()
+			code := []byte(u.code)
+			p.SetCodeAndFile(code, "", u.filepath)
+			output, err := p.Parse(context.Background())
+			assert.NoError(t, err)
+			assert.Equal(t, u.result, output.Modules)
+		})
+	}
+}