Fix deps parser for c++ files

pkg/deps/c.go

@@ -26,7 +26,7 @@ const (
 	StateCImport
 )
 
-// ParserC is a dependency parser for the c programming language.
+// ParserC is a dependency parser for the C programming language.
 // It is not thread safe.
 type ParserC struct {
 	State  StateC

pkg/deps/c_test.go

@@ -10,35 +10,13 @@ import (
 )
 
 func TestParserC_Parse(t *testing.T) {
-	tests := map[string]struct {
-		Filepath string
-		Expected []string
-	}{
-		"c": {
-			Filepath: "testdata/c.c",
-			Expected: []string{
-				"math",
-				"openssl",
-			},
-		},
-		"cpp": {
-			Filepath: "testdata/cpp.cpp",
-			Expected: []string{
-				"iostream",
-				"openssl",
-				"wakatime",
-			},
-		},
-	}
+	parser := deps.ParserC{}
 
-	for name, test := range tests {
-		t.Run(name, func(t *testing.T) {
-			parser := deps.ParserC{}
+	dependencies, err := parser.Parse("testdata/c.c")
+	require.NoError(t, err)
 
-			dependencies, err := parser.Parse(test.Filepath)
-			require.NoError(t, err)
-
-			assert.Equal(t, test.Expected, dependencies)
-		})
-	}
+	assert.Equal(t, []string{
+		"math",
+		"openssl",
+	}, dependencies)
 }

pkg/deps/cpp.go

@@ -0,0 +1,125 @@
+package deps
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"regexp"
+	"strings"
+
+	"github.com/wakatime/wakatime-cli/pkg/heartbeat"
+	"github.com/wakatime/wakatime-cli/pkg/log"
+
+	"github.com/alecthomas/chroma/v2"
+	"github.com/alecthomas/chroma/v2/lexers"
+)
+
+var cppExcludeRegex = regexp.MustCompile(`(?i)^(stdio\.h|iostream|stdlib\.h|string\.h|time\.h)$`)
+
+// StateCPP is a token parsing state.
+type StateCPP int
+
+const (
+	// StateCPPUnknown represents a unknown token parsing state.
+	StateCPPUnknown StateCPP = iota
+	// StateCPPImport means we are in import section during token parsing.
+	StateCPPImport
+)
+
+// ParserCPP is a dependency parser for the C++ programming language.
+// It is not thread safe.
+type ParserCPP struct {
+	State  StateCPP
+	Output []string
+}
+
+// Parse parses dependencies from C++ file content using the C lexer.
+func (p *ParserCPP) Parse(filepath string) ([]string, error) {
+	reader, err := os.Open(filepath) // nolint:gosec
+	if err != nil {
+		return nil, fmt.Errorf("failed to open file %q: %s", filepath, err)
+	}
+
+	defer func() {
+		if err := reader.Close(); err != nil {
+			log.Debugf("failed to close file: %s", err)
+		}
+	}()
+
+	p.init()
+	defer p.init()
+
+	data, err := io.ReadAll(reader)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read from reader: %s", err)
+	}
+
+	l := lexers.Get(heartbeat.LanguageCPP.String())
+	if l == nil {
+		return nil, fmt.Errorf("failed to get lexer for %s", heartbeat.LanguageCPP.String())
+	}
+
+	iter, err := l.Tokenise(nil, string(data))
+	if err != nil {
+		return nil, fmt.Errorf("failed to tokenize file content: %s", err)
+	}
+
+	for _, token := range iter.Tokens() {
+		p.processToken(token)
+	}
+
+	return p.Output, nil
+}
+
+func (p *ParserCPP) append(dep string) {
+	// only consider first part of an import path
+	dep = strings.Split(dep, "/")[0]
+
+	if len(dep) == 0 {
+		return
+	}
+
+	dep = strings.TrimSpace(dep)
+
+	if cppExcludeRegex.MatchString(dep) {
+		return
+	}
+
+	// trim extension
+	dep = strings.TrimSuffix(dep, ".h")
+
+	p.Output = append(p.Output, dep)
+}
+
+func (p *ParserCPP) init() {
+	p.Output = nil
+	p.State = StateCPPUnknown
+}
+
+func (p *ParserCPP) processToken(token chroma.Token) {
+	switch token.Type {
+	case chroma.CommentPreproc:
+		p.processCommentPreproc(token.Value)
+	case chroma.CommentPreprocFile:
+		p.processCommentPreprocFile(token.Value)
+	}
+}
+
+func (p *ParserCPP) processCommentPreproc(value string) {
+	if strings.HasPrefix(strings.TrimSpace(value), "include") {
+		p.State = StateCPPImport
+	}
+}
+
+func (p *ParserCPP) processCommentPreprocFile(value string) {
+	if p.State != StateCPPImport {
+		return
+	}
+
+	if value != "\n" && value != "#" {
+		value = strings.Trim(value, `"<> `)
+		p.append(value)
+	}
+
+	p.State = StateCPPUnknown
+}

pkg/deps/cpp_test.go

@@ -0,0 +1,21 @@
+package deps_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/wakatime/wakatime-cli/pkg/deps"
+)
+
+func TestParserCPP_Parse(t *testing.T) {
+	parser := deps.ParserCPP{}
+
+	dependencies, err := parser.Parse("testdata/cpp.cpp")
+	require.NoError(t, err)
+
+	assert.Equal(t, []string{
+		"openssl",
+		"wakatime",
+	}, dependencies)
+}

pkg/deps/deps.go

@@ -84,8 +84,10 @@ func Detect(filepath string, language heartbeat.Language) ([]string, error) {
 	var parser DependencyParser
 
 	switch language {
-	case heartbeat.LanguageC, heartbeat.LanguageCPP:
+	case heartbeat.LanguageC:
 		parser = &ParserC{}
+	case heartbeat.LanguageCPP:
+		parser = &ParserCPP{}
 	case heartbeat.LanguageCSharp:
 		parser = &ParserCSharp{}
 	case heartbeat.LanguageElm:

pkg/deps/deps_test.go

@@ -162,7 +162,7 @@ func TestDetect(t *testing.T) {
 		"cpp": {
 			Filepath:     "testdata/cpp_minimal.cpp",
 			Language:     heartbeat.LanguageCPP,
-			Dependencies: []string{"iostream"},
+			Dependencies: []string{"wakatime"},
 		},
 		"csharp": {
 			Filepath:     "testdata/csharp_minimal.cs",

pkg/deps/testdata/cpp_minimal.cpp

@@ -1,4 +1,5 @@
 #include <iostream>
+#include "wakatime.h"
 
 int main() {
     std::cout << "Hello World";