Repository: kythe/kythe
Pull request #4084 (2019-09-16 19:45:08)

feat(kzip): add function for merging multiple KzipInfos together
This is useful when extraction results are spread among multiple
kzips and we want to know the combined details.


Patch:
diff --git a/kythe/go/platform/kzip/info/BUILD b/kythe/go/platform/kzip/info/BUILD
index 3e30ec230..e7a8a1e79 100644
--- a/kythe/go/platform/kzip/info/BUILD
+++ b/kythe/go/platform/kzip/info/BUILD
@@ -1,4 +1,4 @@
-load("//tools:build_rules/shims.bzl", "go_library")
+load("//tools:build_rules/shims.bzl", "go_library", "go_test")
 
 package(default_visibility = ["//kythe:default_visibility"])
 
@@ -12,3 +12,13 @@ go_library(
         "//kythe/proto:analysis_go_proto",
     ],
 )
+
+go_test(
+    name = "info_test",
+    srcs = ["info_test.go"],
+    library = ":info",
+    deps = [
+        "//kythe/proto:analysis_go_proto",
+        "@com_github_golang_protobuf//proto:go_default_library",
+    ],
+)
diff --git a/kythe/go/platform/kzip/info/info.go b/kythe/go/platform/kzip/info/info.go
index 4f8f77756..6bb15b034 100644
--- a/kythe/go/platform/kzip/info/info.go
+++ b/kythe/go/platform/kzip/info/info.go
@@ -56,3 +56,34 @@ func KzipInfo(f kzip.File) (*apb.KzipInfo, error) {
 
 	return kzipInfo, nil
 }
+
+// MergeKzipInfo combines the counts from multiple KzipInfos.
+func MergeKzipInfo(infos []*apb.KzipInfo) *apb.KzipInfo {
+	kzipInfo := &apb.KzipInfo{Corpora: make(map[string]*apb.KzipInfo_CorpusInfo)}
+	corpusInfo := func(corpus string) *apb.KzipInfo_CorpusInfo {
+		i := kzipInfo.Corpora[corpus]
+		if i == nil {
+			i = &apb.KzipInfo_CorpusInfo{
+				Files: make(map[string]int32),
+				Units: make(map[string]int32),
+			}
+			kzipInfo.Corpora[corpus] = i
+		}
+		return i
+	}
+
+	for _, i := range infos {
+		kzipInfo.TotalUnits += i.TotalUnits
+		kzipInfo.TotalFiles += i.TotalFiles
+		for corpus, cinfo := range i.Corpora {
+			mergedCorpInfo := corpusInfo(corpus)
+			for lang, fileCount := range cinfo.Files {
+				mergedCorpInfo.Files[lang] += fileCount
+			}
+			for lang, unitCount := range cinfo.Units {
+				mergedCorpInfo.Units[lang] += unitCount
+			}
+		}
+	}
+	return kzipInfo
+}


Test patch:
diff --git a/kythe/go/platform/kzip/info/info_test.go b/kythe/go/platform/kzip/info/info_test.go
new file mode 100644
index 000000000..3d38ce056
--- /dev/null
+++ b/kythe/go/platform/kzip/info/info_test.go
@@ -0,0 +1,69 @@
+/*
+ * Copyright 2019 The Kythe Authors. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package info
+
+import (
+	"testing"
+
+	"github.com/golang/protobuf/proto"
+
+	apb "kythe.io/kythe/proto/analysis_go_proto"
+)
+
+func TestMergeKzipInfo(t *testing.T) {
+	infos := []*apb.KzipInfo{
+		{
+			TotalUnits: 1,
+			TotalFiles: 2,
+			Corpora: map[string]*apb.KzipInfo_CorpusInfo{
+				"corpus1": {
+					Files: map[string]int32{"python": 2},
+					Units: map[string]int32{"python": 1},
+				},
+			},
+		},
+		{
+			TotalUnits: 1,
+			TotalFiles: 2,
+			Corpora: map[string]*apb.KzipInfo_CorpusInfo{
+				"corpus1": {
+					Files: map[string]int32{"go": 2},
+					Units: map[string]int32{"python": 1},
+				},
+			},
+		},
+	}
+
+	want := &apb.KzipInfo{
+		TotalUnits: 2,
+		TotalFiles: 4,
+		Corpora: map[string]*apb.KzipInfo_CorpusInfo{
+			"corpus1": {
+				Files: map[string]int32{
+					"go":     2,
+					"python": 2,
+				},
+				Units: map[string]int32{"python": 2},
+			},
+		},
+	}
+
+	got := MergeKzipInfo(infos)
+	if !proto.Equal(got, want) {
+		t.Errorf("got %v, want %v", got, want)
+	}
+}