From 5ab18313ac06faa3f88773f50a8228488bb406c6 Mon Sep 17 00:00:00 2001
From: Julio Perez <jperez@nvidia.com>
Date: Fri, 6 May 2022 12:18:15 -0400
Subject: [PATCH] add tagset support for select-by-tag

---
 merlin/schema/schema.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/merlin/schema/schema.py b/merlin/schema/schema.py
index 895bd8fd2..f172b9ce6 100644
--- a/merlin/schema/schema.py
+++ b/merlin/schema/schema.py
@@ -285,7 +285,9 @@ def excluding(self, selector) -> "Schema":
     def apply_inverse(self, selector) -> "Schema":
         return self.excluding(selector)
 
-    def select_by_tag(self, tags: Union[Union[str, Tags], List[Union[str, Tags]]]) -> "Schema":
+    def select_by_tag(
+        self, tags: Union[Union[str, Tags], List[Union[str, Tags]], TagSet]
+    ) -> "Schema":
         """Select matching columns from this Schema object using a list of tags
 
         Parameters
@@ -299,8 +301,11 @@ def select_by_tag(self, tags: Union[Union[str, Tags], List[Union[str, Tags]]]) -
             New object containing only the ColumnSchemas of selected columns
 
         """
-        if not isinstance(tags, (list, tuple)):
+        if isinstance(tags, tuple):
+            tags = list(tags)
+        if not isinstance(tags, list):
             tags = [tags]
+        tags = TagSet(tags)
 
         selected_schemas = {}