Completion for untyped accessor methods that has corresponding instance variables

Completes `untyped.foo.` if untyped has an instance variable `@foo`

Author: tompng

lib/repl_type_completor/methods.rb

@@ -5,6 +5,7 @@ module Methods
     OBJECT_SINGLETON_METHODS_METHOD = Object.instance_method(:singleton_methods)
     OBJECT_PRIVATE_METHODS_METHOD = Object.instance_method(:private_methods)
     OBJECT_INSTANCE_VARIABLES_METHOD = Object.instance_method(:instance_variables)
+    OBJECT_INSTANCE_VARIABLE_DEFINED_METHOD = Object.instance_method(:instance_variable_defined?)
     OBJECT_INSTANCE_VARIABLE_GET_METHOD = Object.instance_method(:instance_variable_get)
     OBJECT_CLASS_METHOD = Object.instance_method(:class)
     MODULE_NAME_METHOD = Module.instance_method(:name)

lib/repl_type_completor/type_analyzer.rb

@@ -1179,6 +1179,12 @@ def method_call(receiver, method_name, args, kwargs, block, scope, name_match: t
           end
         end
       end
+
+      if types.empty? && args.empty? && !kwargs && !block
+        t = Types.accessor_method_return_type(receiver, method_name)
+        types << t if t
+      end
+
       scope&.terminate if terminates && breaks.empty?
       Types::UnionType[*types, *breaks]
     end

lib/repl_type_completor/types.rb

@@ -105,6 +105,25 @@ def self.method_return_type(type, method_name)
       UnionType[*types]
     end
 
+    def self.accessor_method_return_type(type, method_name)
+      ivar_name = :"@#{method_name.to_s.gsub(/[!?]\z/, '')}"
+      instances = []
+      type.types.each do |t|
+        case t
+        in SingletonType
+          instances << t.module_or_class
+        in InstanceType
+          instances.concat(t.instances) if t.instances
+        end
+      end
+      instances = instances.sample(OBJECT_TO_TYPE_SAMPLE_SIZE) if instances.size > OBJECT_TO_TYPE_SAMPLE_SIZE
+      objects = []
+      instances.map do |instance|
+        objects << instance.instance_variable_get(ivar_name) if instance.instance_variable_defined?(ivar_name)
+      end
+      union_type_from_objects(objects) if objects
+    end
+
     def self.rbs_methods(type, method_name, args_types, kwargs_type, has_block)
       return [] unless rbs_builder
 
@@ -188,9 +207,10 @@ def self.type_from_object(object)
     end
 
     def self.union_type_from_objects(objects)
-      instanes = objects.size <= OBJECT_TO_TYPE_SAMPLE_SIZE ? objects : objects.sample(OBJECT_TO_TYPE_SAMPLE_SIZE)
-      class_instanes = instanes.group_by { Methods::OBJECT_CLASS_METHOD.bind_call(_1) }
-      UnionType[*class_instanes.map { InstanceType.new _1, nil, _2 }]
+      instances = objects.size <= OBJECT_TO_TYPE_SAMPLE_SIZE ? objects : objects.sample(OBJECT_TO_TYPE_SAMPLE_SIZE)
+      modules, instances = instances.partition { Module === _1 }
+      class_instances = instances.group_by { Methods::OBJECT_CLASS_METHOD.bind_call(_1) }
+      UnionType[*class_instances.map { InstanceType.new _1, nil, _2 }, *modules.uniq.map { SingletonType.new _1 }]
     end
 
     class SingletonType

test/repl_type_completor/test_type_analyze.rb

@@ -286,6 +286,53 @@ def test_method_select
       assert_call('2.times{}.', include: Integer, exclude: Enumerator)
     end
 
+    def test_accessor_method
+      a = Object.new
+      b = Object.new
+      c = Object.new
+      d = Object.new
+      bo = Object.new
+      def a.foo; end
+      def b.foo; end
+      def c.foo; end
+      def d.foo; end
+      def d.bar!; end
+      def d.baz?; end
+      def bo.foo; @foo ||= {}; end
+      bo.foo
+      a.instance_variable_set(:@foo, 1)
+      b.instance_variable_set(:@foo, 'a')
+      c.instance_variable_set(:@foo, nil)
+      d.instance_variable_set(:@bar, :a)
+      d.instance_variable_set(:@baz, [])
+      d.instance_variable_set(:@blah, 1)
+      assert_call('a.foo.', include: Integer, binding:)
+      assert_call('b.foo.', include: String, binding:)
+      assert_call('bo.foo.', include: Hash, binding:)
+      assert_call('[a, b, bo].sample.foo.', include: [Integer, String, Hash], binding:)
+      assert_call('[a, b, c].sample.foo.', include: [Integer, String, NilClass], binding:)
+      assert_call('[a, b, d].sample.foo.', include: [Integer, String], exclude: NilClass, binding:)
+      assert_call('d.bar!.', include: Symbol, binding:)
+      assert_call('d.baz?.', include: Array, binding:)
+      # Even if method does not exist, it might return @blah in method_missing
+      assert_call('d.blah.', include: Integer, binding:)
+    end
+
+    def test_accessor_method_with_class_module
+      AnalyzeTest.instance_variable_set(:@foo, 1)
+      AnalyzeTest.instance_variable_set(:@bar, Symbol)
+      o = Object.new
+      o.instance_variable_set(:@foo, String)
+      o.instance_variable_set(:@bar, Math)
+      assert_call('AnalyzeTest.foo.', include: Integer, binding:)
+      assert_call('AnalyzeTest.bar.all_symbols.', include: Array, binding:)
+      assert_call('o.foo.new.', include: String, binding:)
+      assert_call('o.bar.sin(1).', include: Float, binding:)
+    ensure
+      AnalyzeTest.remove_instance_variable(:@foo)
+      AnalyzeTest.remove_instance_variable(:@bar)
+    end
+
     def test_interface_match_var
       assert_call('([1]+[:a]+["a"]).sample.', include: [Integer, String, Symbol])
     end