Add transform to schema class

README.md

@@ -17,6 +17,9 @@ pip install triad
 
 ## Release History
 
+### 0.3.6
+* Add `transform` to Schema class
+
 ### 0.3.5
 * Change pyarrow and pandas type_safe output to be consistent with pyarrow (None for pd.NaT, nan, etc)
 

tests/collections/test_schema.py

@@ -276,3 +276,20 @@ def test_schema_rename():
     raises(SchemaError, lambda: s.rename(dict(x="b")))
     raises(SchemaError, lambda: s.rename(dict(a="b")))
     raises(SchemaError, lambda: s.rename(dict(a=123)))
+
+
+def test_schema_transform():
+    s = Schema("a:int,b:str,c:bool")
+    assert s.transform("x:str") == "x:str"
+    assert s.transform("*") == s
+    assert s.transform("*~x,y") == s
+    assert s.transform("*,d:str") == "a:int,b:str,c:bool,d:str"
+    assert s.transform("*,d:str - a ") == "b:str,c:bool,d:str"
+    assert s.transform("*,d:str - c,,a ") == "b:str,d:str"
+    assert s.transform("*,d:str ~ c,,a,x ") == "b:str,d:str"
+    assert s.transform("*", {"d": str}, e=str) == "a:int,b:str,c:bool,d:str,e:str"
+    assert s.transform(lambda s: s.fields[0], lambda s: s.fields[2]) == "a:int,c:bool"
+    assert s.transform(lambda s: s - ["b"]) == "a:int,c:bool"
+    raises(SchemaError, lambda: s.transform("**"))
+    raises(SchemaError, lambda: s.transform("*", "*"))
+    raises(SchemaError, lambda: s.transform("*-x"))

triad/__init__.py

@@ -1 +1 @@
-__version__ = "0.3.5"
+__version__ = "0.3.6"

triad/collections/__init__.py

@@ -1,3 +1,4 @@
 # flake8: noqa
 from triad.collections.dict import ParamDict, IndexedOrderedDict
 from triad.collections.schema import Schema
+from triad.collections.fs import FileSystem

triad/collections/schema.py

@@ -21,7 +21,7 @@ class SchemaError(Exception):
     """Exceptions related with construction and modifying schemas
     """
 
-    def __init__(self, message: str):
+    def __init__(self, message: Any):
         super().__init__(message)
 
 
@@ -451,6 +451,51 @@ def rename(self, columns: Dict[str, str], ignore_missing: bool = False) -> "Sche
         ]
         return Schema(pairs)
 
+    def transform(self, *args: Any, **kwargs: Any) -> "Schema":
+        """Transform the current schema to a new schema
+
+        :raises SchemaError: if there is any exception
+        :return: transformed schema
+
+        :Examples:
+        >>> s=Schema("a:int,b:int,c:str")
+        >>> s.transform("x:str") # x:str
+        >>> # add
+        >>> s.transform("*,x:str") # a:int,b:int,c:str,x:str
+        >>> s.transform("*","x:str") # a:int,b:int,c:str,x:str
+        >>> s.transform("*",x=str) # a:int,b:int,c:str,x:str
+        >>> # subtract
+        >>> s.transform("*-c,a") # b:int
+        >>> s.transform("*~c,a,x") # b:int  # ~ means exlcude if exists
+        >>> # callable
+        >>> s.transform(lambda s:s.fields[0]) # a:int
+        >>> s.transform(lambda s:s.fields[0], lambda s:s.fields[2]) # a:int,c:str
+        """
+        try:
+            result = Schema()
+            for a in args:
+                if callable(a):
+                    result += a(self)
+                elif isinstance(a, str):
+                    if "~" in a:
+                        aa = a.split("~", 1)
+                        exclude = True
+                    else:
+                        aa = a.split("-", 1)
+                        exclude = False
+                    s = Schema(aa[0].replace("*", str(self)))
+                    if len(aa) == 2:
+                        cols = [x.strip() for x in aa[1].split(",") if x.strip() != ""]
+                        s = s.exclude(cols) if exclude else s - cols
+                    result += s
+                else:
+                    result += a
+            return result + Schema(kwargs)
+        except SchemaError:
+            raise
+        except Exception as e:
+            raise SchemaError(e)
+
     def _pre_update(self, op: str, need_reindex: bool = True) -> None:
         if op == "__setitem__":
             super()._pre_update(op, need_reindex)