Add Exclude and Extract to the TS-utils test

Author: msullivan

tests/test_ts_utility.py

@@ -6,7 +6,7 @@
 """
 
 import textwrap
-from typing import Literal, Never, NotRequired, Required, TypedDict
+from typing import Literal, Never, NotRequired, Required, TypedDict, Union
 
 import typemap_extensions as typing
 from typemap.type_eval import eval_typing
@@ -50,6 +50,7 @@ class TodoTD(TypedDict):
 
 # Omit<T, Keys>
 # Constructs a type by picking all properties from T and then removing Keys.
+# Note that unlike in TS, our Omit does not depend on Exclude.
 type Omit[T, Keys] = typing.NewProtocol[
     *[
         p
@@ -58,6 +59,27 @@ class TodoTD(TypedDict):
     ]
 ]
 
+# Exclude<T, U>
+# Constructs a type by excluding from T all union members assignable to U.
+type Exclude[T, U] = Union[
+    *[
+        x
+        for x in typing.Iter[typing.FromUnion[T]]
+        if not typing.IsAssignable[x, U]
+    ]
+]
+
+# Extract<T, U>
+# Constructs a type by extracting from T all union members assignable to U.
+type Extract[T, U] = Union[
+    *[
+        x
+        for x in typing.Iter[typing.FromUnion[T]]
+        # Just the inverse of Exclude, really
+        if typing.IsAssignable[x, U]
+    ]
+]
+
 # Partial<T>
 # Constructs a type with all properties of T set to optional (T | None).
 type Partial[T] = typing.NewProtocol[
@@ -156,3 +178,28 @@ def test_td_total_false_inherited():
     assert quals["y"] is Never
     # z defined in total=True child -> no qual
     assert quals["z"] is Never
+
+
+def test_exclude():
+    # bool is assignable to int, so it gets excluded too
+    assert eval_typing(Exclude[str | int | bool, int]) is str
+
+    assert eval_typing(
+        Exclude[Literal["a"] | Literal["b"] | Literal["c"], Literal["a"]]
+    ) == (Literal["b"] | Literal["c"])
+
+    assert eval_typing(Exclude[str | int | float, str | int]) is float
+
+
+def test_extract():
+    # bool is assignable to int, so it gets extracted
+    assert eval_typing(Extract[str | int | bool, int]) == (int | bool)
+
+    assert (
+        eval_typing(
+            Extract[Literal["a"] | Literal["b"] | Literal["c"], Literal["a"]]
+        )
+        == Literal["a"]
+    )
+
+    assert eval_typing(Extract[str | int | float, str | int]) == (str | int)