diff --git a/toy-impl/parse_expr.py b/toy-impl/parse_expr.py index aa0d945..f404440 100644 --- a/toy-impl/parse_expr.py +++ b/toy-impl/parse_expr.py @@ -264,6 +264,9 @@ def parse_list_comprehension(ctx: Context, if node.generators[0].is_async: raise CustomError('async list comprehension is not supported', node) ty = parse_expr(ctx, sym_table, node.generators[0].iter) + if isinstance(ty, TypeVariable) and \ + len(ty.constraints) == 1: + ty = ty.constraints[0] if not isinstance(ty, ListType): raise CustomError(f'unable to iterate over {ty}', node) try: diff --git a/toy-impl/parse_stmt.py b/toy-impl/parse_stmt.py index 075b083..6c60f05 100644 --- a/toy-impl/parse_stmt.py +++ b/toy-impl/parse_stmt.py @@ -159,7 +159,11 @@ def parse_for_stmt(ctx: Context, return_ty: Type, node): ty = parse_expr(ctx, sym_table, node.iter) + if isinstance(ty, TypeVariable) and \ + len(ty.constraints) == 1: + ty = ty.constraints[0] if not isinstance(ty, ListType): + raise CustomError('only iteration over list is supported', node.iter) binding = parse_simple_binding(node.target, ty.params[0]) for key, value in binding.items(): diff --git a/toy-impl/top_level.py b/toy-impl/top_level.py index 86594c7..fbd18b5 100644 --- a/toy-impl/top_level.py +++ b/toy-impl/top_level.py @@ -126,17 +126,12 @@ def parse_type_var(ctx: Context, node): raise CustomError('redefining type variable is not allowed', node) constraints = [] for v in node.value.args[1:]: - if isinstance(v, ast.Constant): - value = v.value - elif isinstance(v, ast.Name): - value = v.id - else: + value, var = parse_type(ctx, v) + if len(var) > 0: raise CustomError( - 'TypeVar constraints must be either string or type name', + 'constraints cannot contain type variables', node) - if value not in ctx.types: - raise CustomError(f'unbounded type {value}', node) - constraints.append(ctx.types[value]) + constraints.append(value) ctx.variables[name] = TypeVariable(node.value.args[0].value, constraints)