fix TypePrinter for Variant and Enum

4 files changed, 69 insertions, 23 deletions

diff --git a/ydb/library/yql/minikql/mkql_type_builder_ut.cpp b/ydb/library/yql/minikql/mkql_type_builder_ut.cpp
index ab270216d8..c8d23b9416 100644
--- a/ydb/library/yql/minikql/mkql_type_builder_ut.cpp
+++ b/ydb/library/yql/minikql/mkql_type_builder_ut.cpp
@@ -120,8 +120,7 @@ private:
             .Build();
 
         auto s = FormatType(t);
-        // todo: fix format: Enum<'BLUE','GREEN','RED'>
-        UNIT_ASSERT_VALUES_EQUAL(s, "Variant<Struct<'BLUE':Void,'GREEN':Void,'RED':Void>>");
+        UNIT_ASSERT_VALUES_EQUAL(s, "Enum<'BLUE','GREEN','RED'>");
     }
 
     void TestResourceTypeFormat() {
@@ -181,7 +180,7 @@ private:
                 .Build();
 
             auto s = FormatType(FunctionTypeInfoBuilder.Variant()->Over(t).Build());
-            UNIT_ASSERT_VALUES_EQUAL(s, "Variant<Struct<'is_ok':Bool,'name':String>>");
+            UNIT_ASSERT_VALUES_EQUAL(s, "Variant<'is_ok':Bool,'name':String>");
         }
         {
             auto t = FunctionTypeInfoBuilder.Tuple(2)
@@ -190,7 +189,7 @@ private:
                 .Build();
 
             auto s = FormatType(FunctionTypeInfoBuilder.Variant()->Over(t).Build());
-            UNIT_ASSERT_VALUES_EQUAL(s, "Variant<Tuple<Bool,String>>");
+            UNIT_ASSERT_VALUES_EQUAL(s, "Variant<Bool,String>");
         }
     }
 
diff --git a/ydb/library/yql/public/udf/udf_type_inspection.h b/ydb/library/yql/public/udf/udf_type_inspection.h
index a940cb2d42..4363d6ad84 100644
--- a/ydb/library/yql/public/udf/udf_type_inspection.h
+++ b/ydb/library/yql/public/udf/udf_type_inspection.h
@@ -85,7 +85,7 @@ using TStubTypeVisitor = TStubTypeVisitor1;
 class TDataTypeInspector: public TStubTypeVisitor
 {
 public:
-    TDataTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TDataTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Data) {
             typeHelper.VisitType(type, this);
         }
@@ -109,7 +109,7 @@ private:
 class TDataAndDecimalTypeInspector: public TStubTypeVisitor
 {
 public:
-    TDataAndDecimalTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TDataAndDecimalTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Data) {
             typeHelper.VisitType(type, this);
         }
@@ -145,7 +145,7 @@ private:
 class TStructTypeInspector: public TStubTypeVisitor
 {
 public:
-    TStructTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TStructTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Struct) {
             typeHelper.VisitType(type, this);
         }
@@ -165,6 +165,8 @@ public:
     }
     const TStringRef& GetMemberName(ui32 i) const { return MembersNames_[i]; }
     const TType* GetMemberType(ui32 i) const { return MembersTypes_[i]; }
+    TStringRef* GetMemberNames() { return MembersNames_.data(); }
+    const TType** GetMemberTypes() { return MembersTypes_.data(); }
 
 private:
     void OnStruct(
@@ -194,7 +196,7 @@ private:
 class TListTypeInspector: public TStubTypeVisitor
 {
 public:
-    TListTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TListTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::List) {
             typeHelper.VisitType(type, this);
         }
@@ -218,7 +220,7 @@ private:
 class TOptionalTypeInspector: public TStubTypeVisitor
 {
 public:
-    TOptionalTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TOptionalTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Optional) {
             typeHelper.VisitType(type, this);
         }
@@ -242,7 +244,7 @@ private:
 class TTupleTypeInspector: public TStubTypeVisitor
 {
 public:
-    TTupleTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TTupleTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Tuple) {
             typeHelper.VisitType(type, this);
         }
@@ -251,6 +253,7 @@ public:
     explicit operator bool() const { return ElementsCount_ != Max<ui32>(); }
     ui32 GetElementsCount() const { return ElementsCount_; }
     const TType* GetElementType(ui32 i) const { return ElementsTypes_[i]; }
+    const TType** GetElementTypes() { return ElementsTypes_.data(); }
 
 private:
     void OnTuple(ui32 elementsCount, const TType** elementsTypes) override {
@@ -273,7 +276,7 @@ private:
 class TDictTypeInspector: public TStubTypeVisitor
 {
 public:
-    TDictTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TDictTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Dict) {
             typeHelper.VisitType(type, this);
         }
@@ -300,7 +303,7 @@ private:
 class TCallableTypeInspector: public TStubTypeVisitor
 {
 public:
-    TCallableTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TCallableTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Callable) {
             typeHelper.VisitType(type, this);
         }
@@ -360,7 +363,7 @@ private:
 class TStreamTypeInspector : public TStubTypeVisitor
 {
 public:
-    TStreamTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TStreamTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Stream) {
             typeHelper.VisitType(type, this);
         }
@@ -384,7 +387,7 @@ private:
 class TVariantTypeInspector : public TStubTypeVisitor
 {
 public:
-    TVariantTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TVariantTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Variant) {
             typeHelper.VisitType(type, this);
         }
@@ -409,7 +412,7 @@ private:
 class TResourceTypeInspector: public TStubTypeVisitor
 {
 public:
-    TResourceTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TResourceTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Resource) {
             typeHelper.VisitType(type, this);
         }
@@ -434,7 +437,7 @@ private:
 class TTaggedTypeInspector: public TStubTypeVisitor
 {
 public:
-    TTaggedTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TTaggedTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Tagged) {
             typeHelper.VisitType(type, this);
         }
@@ -462,7 +465,7 @@ private:
 class TPgTypeInspector: public TStubTypeVisitor
 {
 public:
-    TPgTypeInspector(const ITypeInfoHelper& typeHelper, const TType* type) {
+    TPgTypeInspector(const ITypeInfoHelper1& typeHelper, const TType* type) {
         if (typeHelper.GetTypeKind(type) == ETypeKind::Pg) {
             typeHelper.VisitType(type, this);
         }
diff --git a/ydb/library/yql/public/udf/udf_type_printer.cpp b/ydb/library/yql/public/udf/udf_type_printer.cpp
index fa984fa538..2a0baa3be8 100644
--- a/ydb/library/yql/public/udf/udf_type_printer.cpp
+++ b/ydb/library/yql/public/udf/udf_type_printer.cpp
@@ -1,4 +1,5 @@
 #include "udf_type_printer.h"
+#include "udf_type_inspection.h"
 
 namespace NYql {
 namespace NUdf {
@@ -34,14 +35,17 @@ void TTypePrinter1::OnDataType(TDataTypeId typeId) {
 
 void TTypePrinter1::OnStruct(ui32 membersCount, TStringRef* membersNames, const TType** membersTypes) {
     *Output_ << "Struct<";
+    OutStructPayload(membersCount, membersNames, membersTypes);
+    *Output_ << '>';
+}
+
+void TTypePrinter1::OutStructPayload(ui32 membersCount, TStringRef* membersNames, const TType** membersTypes) {
     for (ui32 i = 0U; i < membersCount; ++i) {
         *Output_ << "'" << std::string_view(membersNames[i]) << "':";
         OutImpl(membersTypes[i]);
         if (i < membersCount - 1U)
             *Output_ << ',';
-
     }
-    *Output_ << '>';
 }
 
 void TTypePrinter1::OnList(const TType* itemType) {
@@ -57,13 +61,16 @@ void TTypePrinter1::OnOptional(const TType* itemType) {
 
 void TTypePrinter1::OnTuple(ui32 elementsCount, const TType** elementsTypes) {
     *Output_ << "Tuple<";
+    OutTuplePayload(elementsCount, elementsTypes);
+    *Output_ << '>';
+}
+
+void TTypePrinter1::OutTuplePayload(ui32 elementsCount, const TType** elementsTypes) {
     for (ui32 i = 0U; i < elementsCount; ++i) {
         OutImpl(elementsTypes[i]);
         if (i < elementsCount - 1U)
             *Output_ << ',';
-
     }
-    *Output_ << '>';
 }
 
 void TTypePrinter1::OnDict(const TType* keyType, const TType* valueType) {
@@ -105,11 +112,42 @@ void TTypePrinter1::OnCallable(const TType* returnType, ui32 argsCount, const TT
 }
 
 void TTypePrinter1::OnVariant(const TType* underlyingType) {
-    *Output_ << "Variant<";
-    OutImpl(underlyingType);
+    switch (TypeHelper1_.GetTypeKind(underlyingType)) {
+    case ETypeKind::Struct: {
+        TStructTypeInspector s(TypeHelper1_, underlyingType);
+        const bool isEnum = std::all_of(s.GetMemberTypes(), s.GetMemberTypes() + s.GetMembersCount(), [this](auto memberType) {
+            return TypeHelper1_.GetTypeKind(memberType) == ETypeKind::Void;
+        });
+
+        if (isEnum) {
+            *Output_ << "Enum<";
+            OutEnumValues(s.GetMembersCount(), s.GetMemberNames());
+        } else {
+            *Output_ << "Variant<";
+            OutStructPayload(s.GetMembersCount(), s.GetMemberNames(), s.GetMemberTypes());
+        }
+        break;
+    }
+    case ETypeKind::Tuple: {
+        TTupleTypeInspector s(TypeHelper1_, underlyingType);
+        *Output_ << "Variant<";
+        OutTuplePayload(s.GetElementsCount(), s.GetElementTypes());
+        break;
+    }
+    default:
+        Y_VERIFY(false, "Unexpected underlying type in Variant");
+    }
     *Output_ << '>';
 }
 
+void TTypePrinter1::OutEnumValues(ui32 membersCount, TStringRef* membersNames) {
+    for (ui32 i = 0U; i < membersCount; ++i) {
+        *Output_ << "'" << std::string_view(membersNames[i]) << '\'';
+        if (i < membersCount - 1U)
+            *Output_ << ',';
+    }
+}
+
 void TTypePrinter1::OnStream(const TType* itemType) {
     *Output_ << "Stream<";
     OutImpl(itemType);
diff --git a/ydb/library/yql/public/udf/udf_type_printer.h b/ydb/library/yql/public/udf/udf_type_printer.h
index a50d1175d8..72b2432010 100644
--- a/ydb/library/yql/public/udf/udf_type_printer.h
+++ b/ydb/library/yql/public/udf/udf_type_printer.h
@@ -28,6 +28,12 @@ protected:
     void OnResourceImpl(TStringRef tag);
     void OnTaggedImpl(const TType* baseType, TStringRef tag);
 
+private:
+    void OutStructPayload(ui32 membersCount, TStringRef* membersNames, const TType** membersTypes);
+    void OutTuplePayload(ui32 elementsCount, const TType** elementsTypes);
+    void OutEnumValues(ui32 membersCount, TStringRef* membersNames);
+
+protected:
     const ITypeInfoHelper1& TypeHelper1_;
     const TType* Type_;
 };