From 594fb475a9d461a75c33876f551e6e71ebed27e3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ernest=20Sury=C5=9B?= <ernest.surys@gmail.com>
Date: Sun, 25 Sep 2022 10:58:06 +0200
Subject: [PATCH] CS Generation Fixes

- Column support
- Template Support
---
 .../UxmlClassGeneration/UxmlClassGenerator.cs | 34 +++++++++++++++++--
 .../Editor/UxmlClassGeneration/UxmlElement.cs | 14 +++++++-
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/Packages/com.quickeye.ui-toolkit-plus/Editor/UxmlClassGeneration/UxmlClassGenerator.cs b/Packages/com.quickeye.ui-toolkit-plus/Editor/UxmlClassGeneration/UxmlClassGenerator.cs
index b647fc2..c38c47e 100644
--- a/Packages/com.quickeye.ui-toolkit-plus/Editor/UxmlClassGeneration/UxmlClassGenerator.cs
+++ b/Packages/com.quickeye.ui-toolkit-plus/Editor/UxmlClassGeneration/UxmlClassGenerator.cs
@@ -9,7 +9,15 @@ namespace QuickEye.UIToolkit.Editor
     internal static class UxmlClassGenerator
     {
         private const string GenerateCsMenuItemName = "CONTEXT/VisualTreeAsset/Generate C# Class";
-        
+
+        private static readonly string[] IgnoredTagFullnames =
+        {
+            "UnityEngine.UIElements.Template",
+            "Style"
+        };
+
+        private const string ColumnFullName = "UnityEngine.UIElements.Column";
+
         [MenuItem(GenerateCsMenuItemName)]
         private static void GenerateGenCs(MenuCommand command)
         {
@@ -30,7 +38,10 @@ public static void GenerateGenCs(VisualTreeAsset asset)
             var settings = CodeGenSettings.FromUxml(uxml);
             if (UxmlParser.TryGetElementsWithName(uxml, out var elements))
             {
-                GenerateScript(asset.name, elements, GetGenCsFilePath(path), settings);
+                var validElements = elements
+                    .Where(e => !IgnoredTagFullnames.Contains(e.FullyQualifiedTypeName))
+                    .ToArray();
+                GenerateScript(asset.name, validElements, GetGenCsFilePath(path), settings);
             }
         }
 
@@ -48,9 +59,28 @@ private static string GetFieldAssigment(UxmlElement element, CodeGenSettings set
             var name = element.NameAttribute;
             var varName = settings.FieldPrefix + CodeGeneration.UssNameToVariableName(name);
 
+            if (element.FullyQualifiedTypeName == ColumnFullName)
+            {
+                return GetFieldAssigmentForColumn(element, settings, varName, name);
+            }
+
             return $"{varName} = root.Q<{type}>(\"{name}\");";
         }
 
+        private static string GetFieldAssigmentForColumn(UxmlElement element, CodeGenSettings settings, string varName,
+            string name)
+        {
+            var multiColumnElement = element.XElement.Parent?.Parent?.ToUxmlElement();
+
+            if (string.IsNullOrEmpty(multiColumnElement?.NameAttribute))
+                return $"// Could not find \"{name}\" MultiColumn parent with a name.";
+
+            var multiColumnEleVarName = settings.FieldPrefix +
+                                        CodeGeneration.UssNameToVariableName(multiColumnElement
+                                            .NameAttribute);
+            return $"{varName} = {multiColumnEleVarName}.columns[\"{name}\"];";
+        }
+
         private static void GenerateScript(string scriptName, UxmlElement[] uxmlElements, string path,
             CodeGenSettings settings)
         {
diff --git a/Packages/com.quickeye.ui-toolkit-plus/Editor/UxmlClassGeneration/UxmlElement.cs b/Packages/com.quickeye.ui-toolkit-plus/Editor/UxmlClassGeneration/UxmlElement.cs
index 89105c6..25333a0 100644
--- a/Packages/com.quickeye.ui-toolkit-plus/Editor/UxmlClassGeneration/UxmlElement.cs
+++ b/Packages/com.quickeye.ui-toolkit-plus/Editor/UxmlClassGeneration/UxmlElement.cs
@@ -5,14 +5,16 @@ namespace QuickEye.UIToolkit.Editor
 {
     internal class UxmlElement
     {
+        public readonly XElement XElement;
         public readonly string Namespace;
         public readonly string TypeName;
         public readonly string NameAttribute;
         public string FullyQualifiedTypeName => $"{Namespace}.{TypeName}";
         public bool IsUnityEngineType => Namespace == "UnityEngine.UIElements";
-
+        
         public UxmlElement(XElement xElement)
         {
+            XElement = xElement;
             var localName = xElement.Name.LocalName;
 
             Namespace = xElement.Name.NamespaceName != XNamespace.None
@@ -22,6 +24,16 @@ public UxmlElement(XElement xElement)
                 ? localName.Split('.').Last()
                 : localName;
             NameAttribute = xElement.Attribute("name")?.Value;
+            if (FullyQualifiedTypeName == "UnityEngine.UIElements.Instance")
+                TypeName = "TemplateContainer";
+        }
+    }
+
+    internal static class XElementExtensions
+    {
+        public static UxmlElement ToUxmlElement(this XElement xElement)
+        {
+            return new UxmlElement(xElement);
         }
     }
 }
\ No newline at end of file