Create a gist now

Instantly share code, notes, and snippets.

require 'set'
class InferType
def self.run(&block)
infer_type = InferType.new
infer_type.start
block.call
ensure
infer_type.finish
end
def initialize
@records = {}
end
def start
@trace = TracePoint.new(:call, &method(:on_call))
@trace.enable
end
def finish
@trace.disable
report
end
private
def on_call(tp)
meth = find_method(tp)
if meth
meth.parameters.each do |arg_type, arg_name|
if arg_name
sig = build_signature(tp, meth)
@records[sig] ||= {}
@records[sig][arg_name] ||= Set.new
@records[sig][arg_name].add(tp.binding.local_variable_get(arg_name).class)
end
end
end
end
def find_method(tp)
tp.defined_class.instance_method(tp.method_id)
rescue NameError
nil
end
def build_signature(tp, meth)
sig =
if tp.defined_class.singleton_class?
attached = tp.defined_class.inspect.slice(/#<Class:(.+)>/, 1)
"#{attached}.#{tp.method_id}"
else
"#{tp.defined_class}##{tp.method_id}"
end
loc = meth.source_location
if loc
"#{sig} #{loc[0]}:#{loc[1]}"
else
sig
end
end
def report
re = Regexp.new(ENV.fetch('INFER_TYPE_TARGET', '.'))
@records.each do |signature, method_info|
if re === signature
puts signature
method_info.each do |arg_name, klasses|
t = infer_type(klasses.to_a)
puts " @param #{arg_name} [#{t}]"
end
end
end
end
def infer_type(klasses)
if klasses.size == 1 && klasses.member?(NilClass)
return 'nil'
end
nullable = klasses.member?(NilClass)
klasses.delete(NilClass)
if klasses.size == 2 && klasses.member?(TrueClass) && klasses.member?(FalseClass)
return annotate_nullable('Boolean', nullable)
end
common = klasses.drop(1).inject(klasses[0], &method(:common_type))
annotate_nullable(common, nullable)
end
def annotate_nullable(t, nullable)
if nullable
"#{t}, nil"
else
t
end
end
def common_type(t1, t2)
top = BasicObject
parents(t1).zip(parents(t2)) do |x, y|
unless x.equal?(y)
return top
end
top = x
end
top
end
def parents(t)
klasses = []
until t.equal?(BasicObject)
klasses << t
t = t.superclass
end
klasses << BasicObject
klasses.reverse
end
end
if $0 == __FILE__
path = ARGV.shift
InferType.run { load path }
end
# require_relative '../infer_type'
# RSpec.configure do |config|
# config.before(:suite) { @infer_type = InferType.new; @infer_type.start }
# config.after(:suite) { @infer_type.finish }
# end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment