YT: Allow to define enum with custom string represenation

2 files changed, 30 insertions, 0 deletions

diff --git a/library/cpp/yt/misc/enum-inl.h b/library/cpp/yt/misc/enum-inl.h
index 86515e8a4e3..18a78621a9c 100644
--- a/library/cpp/yt/misc/enum-inl.h
+++ b/library/cpp/yt/misc/enum-inl.h
@@ -140,6 +140,22 @@ static constexpr bool AreValuesDistinct(const TValues& values)
     )(item)
 
 #define ENUM__GET_DOMAIN_NAMES_ITEM_SEQ(seq) \
+    PP_IF( \
+       ENUM__ITEM_SEQ_HAS_DOMAIN_NAME(seq), \
+       ENUM__GET_DOMAIN_NAMES_ITEM_SEQ_CUSTOM, \
+       ENUM__GET_DOMAIN_NAMES_ITEM_SEQ_AUTO \
+    )(seq)
+
+#define ENUM__ITEM_SEQ_HAS_DOMAIN_NAME(seq) \
+    PP_CONCAT(ENUM__ITEM_SEQ_HAS_DOMAIN_NAME_, PP_COUNT(seq))
+
+#define ENUM__ITEM_SEQ_HAS_DOMAIN_NAME_2 PP_FALSE
+#define ENUM__ITEM_SEQ_HAS_DOMAIN_NAME_3 PP_TRUE
+
+#define ENUM__GET_DOMAIN_NAMES_ITEM_SEQ_CUSTOM(seq) \
+    TStringBuf(PP_ELEMENT(seq, 2)),
+
+#define ENUM__GET_DOMAIN_NAMES_ITEM_SEQ_AUTO(seq) \
     ENUM__GET_DOMAIN_NAMES_ITEM_ATOMIC(PP_ELEMENT(seq, 0))
 
 #define ENUM__GET_DOMAIN_NAMES_ITEM_ATOMIC(item) \
diff --git a/library/cpp/yt/misc/unittests/enum_ut.cpp b/library/cpp/yt/misc/unittests/enum_ut.cpp
index 6af11b04dc0..2a024be228c 100644
--- a/library/cpp/yt/misc/unittests/enum_ut.cpp
+++ b/library/cpp/yt/misc/unittests/enum_ut.cpp
@@ -37,6 +37,11 @@ DEFINE_AMBIGUOUS_ENUM_WITH_UNDERLYING_TYPE(EMultipleNames, int,
     ((D2)(100))
 );
 
+DEFINE_ENUM(ECustomString,
+    ((A) (1) ("1_a"))
+    ((B) (2) ("1_b"))
+);
+
 ////////////////////////////////////////////////////////////////////////////////
 
 template <class T, size_t N>
@@ -243,6 +248,15 @@ TEST(TEnumTest, MultipleNames)
     EXPECT_EQ("D1", ToString(EMultipleNames::D2));
 }
 
+TEST(TEnumTest, CustomString)
+{
+    EXPECT_EQ(ECustomString::A, TEnumTraits<ECustomString>::FromString("1_a"));
+    EXPECT_EQ(ECustomString::B, TEnumTraits<ECustomString>::FromString("1_b"));
+
+    EXPECT_EQ("1_a", ToString(ECustomString::A));
+    EXPECT_EQ("1_b", ToString(ECustomString::B));
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace