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 + #include "absl/strings/ascii.h" #include "absl/strings/str_split.h" #include "absl/strings/string_view.h" @@ -40,11 +42,63 @@ std::vector SplitNameAndMetadata( return parts; } +// Use comma as separate to split input metadata. However, treat comma inside +// ""/''/[]/{}/() pairs as normal characters. +std::vector SplitPairs(absl::string_view metadata) { + std::vector key_value_pairs; + std::stack 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> ParseMetadata( absl::string_view metadata) { std::vector> key_values; - for (absl::string_view pair : - absl::StrSplit(metadata, ',', absl::SkipWhitespace())) { + for (absl::string_view pair : SplitPairs(metadata)) { std::vector 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