diff --git a/lib/typeprof/core/ast.rb b/lib/typeprof/core/ast.rb index 54f05878..ee91c0bc 100644 --- a/lib/typeprof/core/ast.rb +++ b/lib/typeprof/core/ast.rb @@ -473,6 +473,28 @@ def self.create_rbs_member(raw_decl, lenv) end end + def self.parse_type_var_comment(raw_node, lenv) + comments = lenv.file_context.comments + return nil unless comments + node_line = raw_node.location.start_line + idx = comments.bsearch_index { |c| c.location.start_line >= node_line } + idx = (idx || comments.size) - 1 + return nil if idx < 0 + comment = comments[idx] + return nil unless comment.location.start_line == node_line - 1 + text = comment.location.slice + if text =~ /\A#\s*@type\s+var\s+(\w+)\s*:\s*(.+)\z/ + var_name = $1.to_sym + type_str = $2 + rbs_type = RBS::Parser.parse_type(type_str) + return nil unless rbs_type + rbs_type_node = AST.create_rbs_type(rbs_type, lenv) + [var_name, rbs_type_node] + end + rescue RBS::ParsingError + nil + end + def self.create_rbs_func_type(raw_decl, raw_type_params, raw_block, lenv) SigFuncType.new(raw_decl, raw_type_params, raw_block, lenv) end diff --git a/lib/typeprof/core/ast/misc.rb b/lib/typeprof/core/ast/misc.rb index 52db028d..c946740e 100644 --- a/lib/typeprof/core/ast/misc.rb +++ b/lib/typeprof/core/ast/misc.rb @@ -12,18 +12,40 @@ def initialize(raw_node, lenv, use_result) DummyNilNode.new(TypeProf::CodeRange.new(last, last), lenv) end end + @type_var_assertions = {} + stmts.each_with_index do |n, i| + next unless n + result = AST.parse_type_var_comment(n, lenv) + if result + @type_var_assertions[i] = result + end + end end attr_reader :stmts def subnodes = { stmts: } + def define0(genv) + @type_var_assertions.each_value do |_var_name, rbs_type_node| + rbs_type_node.define(genv) + end + super(genv) + end + + def undefine0(genv) + super(genv) + @type_var_assertions.each_value do |_var_name, rbs_type_node| + rbs_type_node.undefine(genv) if rbs_type_node.static_ret + end + end + def install0(genv) ret = nil post_stmts = [] - @stmts.each do |stmt| + @stmts.each_with_index do |stmt, i| next if stmt.nil? if stmt.is_a?(PostExecutionNode) @@ -31,6 +53,12 @@ def install0(genv) next end + if (assertion = @type_var_assertions[i]) + var_name, rbs_type_node = assertion + box = @changes.add_type_read_box(genv, rbs_type_node) + @lenv.set_var(var_name, box.ret) + end + ret = stmt.install(genv) end diff --git a/scenario/misc/type_var_comment.rb b/scenario/misc/type_var_comment.rb new file mode 100644 index 00000000..fb18b963 --- /dev/null +++ b/scenario/misc/type_var_comment.rb @@ -0,0 +1,11 @@ +## update: test.rb +def foo(x) + a = x + # @type var a: Integer + a + 1 +end + +## assert +class Object + def foo: (untyped) -> Integer +end