From 8b419dee7bed8b3e267cf698136cfa2f42ded303 Mon Sep 17 00:00:00 2001
From: candiduslynx <candiduslynx@gmail.com>
Date: Wed, 18 Oct 2023 10:20:00 +0300
Subject: [PATCH] fix key construction in `ExtractGoComments`

---
 comment_extractor.go | 42 ++++++++++++++++++++++++++----------------
 go.mod               |  1 +
 go.sum               |  2 ++
 reflect_test.go      |  2 +-
 4 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/comment_extractor.go b/comment_extractor.go
index e157837..1124900 100644
--- a/comment_extractor.go
+++ b/comment_extractor.go
@@ -2,15 +2,16 @@ package jsonschema
 
 import (
 	"fmt"
+	"go/ast"
+	"go/doc"
+	"go/parser"
+	"go/token"
 	"io/fs"
 	gopath "path"
 	"path/filepath"
 	"strings"
 
-	"go/ast"
-	"go/doc"
-	"go/parser"
-	"go/token"
+	"golang.org/x/exp/maps"
 )
 
 // ExtractGoComments will read all the go files contained in the provided path,
@@ -24,24 +25,33 @@ import (
 //
 // When parsing type comments, we use the `go/doc`'s Synopsis method to extract the first phrase
 // only. Field comments, which tend to be much shorter, will include everything.
-func ExtractGoComments(base, path string, commentMap map[string]string) error {
+func ExtractGoComments(base, rootPath string, commentMap map[string]string) error {
+	root, err := filepath.Abs(rootPath)
+	if err != nil {
+		return err
+	}
+
 	fset := token.NewFileSet()
 	dict := make(map[string][]*ast.Package)
-	err := filepath.Walk(path, func(path string, info fs.FileInfo, err error) error {
+	err = filepath.Walk(root, func(path string, info fs.FileInfo, err error) error {
 		if err != nil {
 			return err
 		}
-		if info.IsDir() {
-			d, err := parser.ParseDir(fset, path, nil, parser.ParseComments)
-			if err != nil {
-				return err
-			}
-			for _, v := range d {
-				// paths may have multiple packages, like for tests
-				k := gopath.Join(base, path)
-				dict[k] = append(dict[k], v)
-			}
+		if !info.IsDir() {
+			return nil
 		}
+
+		d, err := parser.ParseDir(fset, path, nil, parser.ParseComments)
+		if err != nil {
+			return err
+		}
+
+		// key should consist of a base path + the path relative to the root
+		// so that we can treat rootPath of "../" properly
+		k := gopath.Join(base, strings.TrimPrefix(path, root))
+		// paths may have multiple packages, like for tests
+		dict[k] = append(dict[k], maps.Values(d)...)
+
 		return nil
 	})
 	if err != nil {
diff --git a/go.mod b/go.mod
index 49c0e35..14447f3 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.18
 require (
 	github.com/stretchr/testify v1.8.1
 	github.com/wk8/go-ordered-map/v2 v2.1.8
+	golang.org/x/exp v0.0.0-20231006140011-7918f672742d
 )
 
 require (
diff --git a/go.sum b/go.sum
index 025db2e..0937aeb 100644
--- a/go.sum
+++ b/go.sum
@@ -19,6 +19,8 @@ github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKs
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
 github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
+golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
+golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/reflect_test.go b/reflect_test.go
index 94b6018..f166704 100644
--- a/reflect_test.go
+++ b/reflect_test.go
@@ -485,7 +485,7 @@ func TestSchemaGeneration(t *testing.T) {
 func prepareCommentReflector(t *testing.T) *Reflector {
 	t.Helper()
 	r := new(Reflector)
-	err := r.AddGoComments("github.com/invopop/jsonschema", "./examples")
+	err := r.AddGoComments("github.com/invopop/jsonschema/examples", "./examples")
 	require.NoError(t, err, "did not expect error while adding comments")
 	return r
 }