experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

allow comma to appear in the key or value of trace arguments. as long as it is quoted or enclosed in bracket/braces etc.

Browse Source 
PiperOrigin-RevId: 290369090
Change-Id: Id7c60f236f0814d0d51c7cbe8e4d9875408905cb
This commit is contained in:
A. Unique TensorFlower 2020-01-17 17:03:48 -08:00 committed by TensorFlower Gardener
parent 6fff5dea63
commit 659ff38c2f
2 changed files with 80 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
tensorflow/core/profiler/internal/parse_annotation.cc
View File

@ -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) {

24
tensorflow/core/profiler/internal/parse_annotation_test.cc
View File

@ -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