From 659ff38c2f138ec6470996551a557a0f0ee872fe Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 Jan 2020 17:03:48 -0800
Subject: [PATCH] allow comma to appear in the key or value of trace arguments.
as long as it is quoted or enclosed in bracket/braces etc.
PiperOrigin-RevId: 290369090
Change-Id: Id7c60f236f0814d0d51c7cbe8e4d9875408905cb
---
.../profiler/internal/parse_annotation.cc | 58 ++++++++++++++++++-
.../internal/parse_annotation_test.cc | 24 ++++++++
2 files changed, 80 insertions(+), 2 deletions(-)
diff --git a/tensorflow/core/profiler/internal/parse_annotation.cc b/tensorflow/core/profiler/internal/parse_annotation.cc
index 8a5d21c79f5..2a3fa3f8454 100644
--- a/tensorflow/core/profiler/internal/parse_annotation.cc
+++ b/tensorflow/core/profiler/internal/parse_annotation.cc
@@ -14,6 +14,8 @@ limitations under the License.
==============================================================================*/
#include "tensorflow/core/profiler/internal/parse_annotation.h"
+#include <stack>
+
#include "absl/strings/ascii.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
@@ -40,11 +42,63 @@ std::vector<absl::string_view> SplitNameAndMetadata(
return parts;
}
+// Use comma as separate to split input metadata. However, treat comma inside
+// ""/''/[]/{}/() pairs as normal characters.
+std::vector<absl::string_view> SplitPairs(absl::string_view metadata) {
+ std::vector<absl::string_view> key_value_pairs;
+ std::stack<char> quotes;
+ int start = 0, end = 0;
+ for (; end < metadata.size(); ++end) {
+ char ch = metadata[end];
+ switch (ch) {
+ case '\"':
+ case '\'':
+ if (quotes.empty() || quotes.top() != ch) {
+ quotes.push(ch);
+ } else {
+ quotes.pop();
+ }
+ break;
+ case '{':
+ case '(':
+ case '[':
+ quotes.push(ch);
+ break;
+ case '}':
+ if (!quotes.empty() && quotes.top() == '{') {
+ quotes.pop();
+ }
+ break;
+ case ')':
+ if (!quotes.empty() && quotes.top() == '(') {
+ quotes.pop();
+ }
+ break;
+ case ']':
+ if (!quotes.empty() && quotes.top() == '[') {
+ quotes.pop();
+ }
+ break;
+ case ',':
+ if (quotes.empty()) {
+ if (end - start > 1) {
+ key_value_pairs.emplace_back(metadata.data() + start, end - start);
+ }
+ start = end + 1; // Skip the current ','.
+ }
+ break;
+ }
+ }
+ if (end - start > 1) {
+ key_value_pairs.emplace_back(metadata.data() + start, end - start);
+ }
+ return key_value_pairs;
+}
+
std::vector<std::pair<absl::string_view, absl::string_view>> ParseMetadata(
absl::string_view metadata) {
std::vector<std::pair<absl::string_view, absl::string_view>> key_values;
- for (absl::string_view pair :
- absl::StrSplit(metadata, ',', absl::SkipWhitespace())) {
+ for (absl::string_view pair : SplitPairs(metadata)) {
std::vector<absl::string_view> parts =
absl::StrSplit(pair, absl::MaxSplits('=', 1));
if (parts.size() == 2) {
diff --git a/tensorflow/core/profiler/internal/parse_annotation_test.cc b/tensorflow/core/profiler/internal/parse_annotation_test.cc
index 65d4ed7d7c3..4d4a2d5ea95 100644
--- a/tensorflow/core/profiler/internal/parse_annotation_test.cc
+++ b/tensorflow/core/profiler/internal/parse_annotation_test.cc
@@ -123,6 +123,30 @@ TEST(ParseAnnotationTest, ExtraMetadataSeparatorTest) {
EXPECT_TRUE(annotation.metadata.empty());
}
+TEST(ParseAnnotationTest, QuotedMetadata) {
+ Annotation annotation = ParseAnnotation(
+ "name#k1=(v11,v12),k2=[v21,v22,v23],k3={v31,v32}, k4=\"v41,v42\","
+ "(k51,k52)='v51,v52'#");
+ EXPECT_EQ(annotation.metadata.at(0).key, "k1");
+ EXPECT_EQ(annotation.metadata.at(0).value, "(v11,v12)");
+ EXPECT_EQ(annotation.metadata.at(1).key, "k2");
+ EXPECT_EQ(annotation.metadata.at(1).value, "[v21,v22,v23]");
+ EXPECT_EQ(annotation.metadata.at(2).key, "k3");
+ EXPECT_EQ(annotation.metadata.at(2).value, "{v31,v32}");
+ EXPECT_EQ(annotation.metadata.at(3).key, "k4");
+ EXPECT_EQ(annotation.metadata.at(3).value, "\"v41,v42\"");
+ EXPECT_EQ(annotation.metadata.at(4).key, "(k51,k52)");
+ EXPECT_EQ(annotation.metadata.at(4).value, "'v51,v52'");
+}
+
+// Make sure unmatched quotes don't die.
+TEST(ParseAnnotationTest, UnmatchedQuotedMetadata) {
+ Annotation annotation = ParseAnnotation("name#k1=v1,k2=(v2,k3=v3#");
+ EXPECT_EQ(annotation.metadata.at(0).key, "k1");
+ EXPECT_EQ(annotation.metadata.at(0).value, "v1");
+ EXPECT_EQ(annotation.metadata.at(1).key, "k2");
+ EXPECT_EQ(annotation.metadata.at(1).value, "(v2,k3=v3");
+}
} // namespace
} // namespace profiler
} // namespace tensorflow