Add functionality for union types type substitution

Author: alexisthedev

src/ir/type_utils.py

@@ -1101,10 +1101,8 @@ class B : A<String>()
                     if not _update_type_var_map(type_var_map, t_var, t_arg1):
                         return {}
                     continue
-                is_parameterized = isinstance(t_var.bound,
-                                              tp.ParameterizedType)
-                is_parameterized2 = isinstance(t_arg1,
-                                               tp.ParameterizedType)
+                is_parameterized = t_var.bound.is_combound()
+                is_parameterized2 = t_arg1.is_combound()
                 if is_parameterized and is_parameterized2:
                     res = unify_types(t_arg1, t_var.bound, factory)
                     if not res or any(

src/ir/types.py

@@ -100,6 +100,9 @@ def is_type_var(self):
     def is_wildcard(self):
         return False
 
+    def is_combound(self):
+        return False
+
     def is_parameterized(self):
         return False
 
@@ -351,7 +354,7 @@ def get_type_variables(self, factory):
             return self.bound.get_type_variables(factory)
         elif self.bound.is_type_var():
             return {self.bound: {self.bound.get_bound_rec(factory)}}
-        elif self.bound.is_parameterized():
+        elif self.bound.is_combound():
             return self.bound.get_type_variables(factory)
         else:
             return {}
@@ -536,7 +539,7 @@ def _to_type_variable_free(t: Type, t_param, factory) -> Type:
             )
         )
         return WildCardType(bound, variance)
-    elif t.is_parameterized():
+    elif t.is_combound():
         return t.to_type_variable_free(factory)
     else:
         return t
@@ -592,6 +595,9 @@ def __init__(self, t_constructor: TypeConstructor,
         # XXX revisit
         self.supertypes = copy(self.t_constructor.supertypes)
 
+    def is_combound(self):
+        return True
+
     def is_parameterized(self):
         return True
 
@@ -661,12 +667,11 @@ def get_type_variables(self, factory) -> Dict[TypeParameter, Set[Type]]:
         # This function actually returns a dict of the enclosing type variables
         # along with the set of their bounds.
         type_vars = defaultdict(set)
-        for i, t_arg in enumerate(self.type_args):
-            t_arg = t_arg
+        for t_arg in self.type_args:
             if t_arg.is_type_var():
                 type_vars[t_arg].add(
                     t_arg.get_bound_rec(factory))
-            elif t_arg.is_parameterized() or t_arg.is_wildcard():
+            elif t_arg.is_combound() or t_arg.is_wildcard():
                 for k, v in t_arg.get_type_variables(factory).items():
                     type_vars[k].update(v)
             else:

src/ir/typescript_types.py

@@ -1,3 +1,5 @@
+from collections import defaultdict
+
 import src.ir.ast as ast
 import src.ir.typescript_ast as ts_ast
 import src.ir.builtins as bt
@@ -96,7 +98,7 @@ def update_add_node_to_parent(self):
 
     def get_dynamic_types(self, gen_object):
         return [
-            #union_types.get_union_type(gen_object),
+            union_types.get_union_type(gen_object),
         ]
 
     def get_constant_candidates(self, constants):
@@ -430,6 +432,9 @@ def __init__(self, types, name="UnionType", primitive=False):
     def get_types(self):
         return self.types
 
+    def is_combound(self):
+        return True
+
     @two_way_subtyping
     def is_subtype(self, other):
         if isinstance(other, UnionType):
@@ -439,9 +444,53 @@ def is_subtype(self, other):
     def dynamic_subtyping(self, other):
         return other in set(self.types)
 
+    def substitute_type_args(self, type_map,
+                             cond=lambda t: t.has_type_variables()):
+        new_types = []
+        for t in self.types:
+            new_t = (t.substitute_type_args(type_map, cond)
+                     if t.has_type_variables()
+                     else t)
+            new_types.append(new_t)
+        return UnionType(new_types)
+
     def has_type_variables(self):
         return any(t.has_type_variables() for t in self.types)
 
+    def get_type_variables(self, factory):
+        # This function actually returns a dict of the enclosing type variables
+        # along with the set of their bounds.
+        type_vars = defaultdict(set)
+        for t in self.types:
+            if t.is_type_var():
+                type_vars[t].add(
+                    t.get_bound_rec(factory))
+            elif t.is_combound() or t.is_wildcard():
+                for k, v in t.get_type_variables(factory).items():
+                    type_vars[k].update(v)
+            else:
+                continue
+        return type_vars
+
+    def to_variance_free(self, type_var_map=None):
+        new_types = []
+        for t in self.types:
+            new_types.append(t.to_variance_free(type_var_map)
+                             if t.is_combound()
+                             else t)
+        return UnionType(new_types)
+
+    def to_type_variable_free(self, factory):
+        # We translate a union type that contains
+        # type variables into a parameterized type that is
+        # type variable free.
+        new_types = []
+        for t in self.types:
+            new_types.append(t.to_type_variable_free(factory)
+                             if t.is_combound()
+                             else t)
+        return UnionType(new_types)
+
     def get_name(self):
         return self.name
 

src/translators/typescript.py

@@ -78,7 +78,7 @@ def get_incorrect_filename():
         return TypeScriptTranslator.incorrect_filename
 
     def get_union(self, utype):
-        return " | ".join([self.type_arg2str(t, True) for t in utype.types])
+        return " | ".join([self.get_type_name(t, True) for t in utype.types])
 
     def type_arg2str(self, t_arg, from_union=False):
         # TypeScript does not have a Wildcard type
@@ -108,7 +108,7 @@ def get_type_name(self, t, from_union=False):
                 self.type_arg2str(ret_type)
             )
             if from_union:
-                return "("+res+")"
+                return "(" + res + ")"
             return res
 
         return "{}<{}>".format(t.name, ", ".join([self.type_arg2str(ta)

tests/test_typescript.py

@@ -57,3 +57,19 @@ def test_union_type_param():
     assert not union2.is_subtype(union1)
     assert not union1.is_subtype(t_param)
     assert not t_param.is_subtype(union1)
+
+
+def test_union_type_substitution():
+    type_param1 = tp.TypeParameter("T1")
+    type_param2 = tp.TypeParameter("T2")
+    type_param3 = tp.TypeParameter("T3")
+    type_param4 = tp.TypeParameter("T4")
+
+    foo = tp.TypeConstructor("Foo", [type_param1, type_param2])
+    foo_p = foo.new([tst.NumberType(), type_param3])
+
+    union = tst.UnionType([tst.StringLiteralType("bar"), foo_p])
+    ptype = tp.substitute_type(union, {type_param3: type_param4})
+
+    assert ptype.types[1].type_args[0] == tst.NumberType()
+    assert ptype.types[1].type_args[1] == type_param4