simplify list return type check

vyperlang · Jan 3, 2020 · 6096c45 · 6096c45
1 parent 4ac5ef5
commit 6096c45
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/vyper/parser/stmt.py b/vyper/parser/stmt.py
@@ -1,5 +1,3 @@
-import re
-
 from vyper import ast
 from vyper.ast_utils import (
     ast_to_dict,
@@ -877,13 +875,11 @@ def parse_return(self):
                 raise Exception(f"Invalid location: {sub.location}")
 
         elif isinstance(sub.typ, ListType):
-            sub_base_type = re.split(r'\(|\[', str(sub.typ.subtype))[0]
-            ret_base_type = re.split(r'\(|\[', str(self.context.return_type.subtype))[0]
             loop_memory_position = self.context.new_placeholder(typ=BaseType('uint256'))
-            if sub_base_type != ret_base_type:
+            if sub.typ.subtype != self.context.return_type.subtype:
                 raise TypeMismatchException(
-                    f"List return type {sub_base_type} does not match specified "
-                    f"return type, expecting {ret_base_type}",
+                    f"List return type {sub.typ.subtype} does not match specified "
+                    f"return type, expecting {self.context.return_type.subtype}",
                     self.stmt
                 )
             elif sub.location == "memory" and sub.value != "multi":