Support instance self class

lib/steep/type_construction.rb

@@ -436,16 +436,39 @@ def for_class(node)
     def for_sclass(node, type)
       annots = source.annotations(block: node, factory: checker.factory, current_module: current_namespace)
 
-      type = module_context.instance_type if type.is_a?(AST::Types::Self)
+      instance_type = if type.is_a?(AST::Types::Self)
+                        context.self_type
+                      else
+                        type
+                      end
 
-      instance_type, module_type = case
-                                   when checker.factory.class_name?(type.name)
-                                     [type.to_class(constructor: nil), AST::Builtin::Class.instance_type]
-                                   when checker.factory.module_name?(type.name)
-                                     [type.to_module, AST::Builtin::Module.instance_type]
-                                   else
-                                     raise "Is #{type.name} a class or module???"
-                                   end
+      module_type = case instance_type
+                    when AST::Types::Name::Class
+                      AST::Builtin::Class.instance_type
+                    when AST::Types::Name::Module
+                      AST::Builtin::Module.instance_type
+                    when AST::Types::Name::Instance
+                      instance_type.to_class(constructor: nil)
+                    else
+                      raise "Unexpected type for sclass node: #{type}"
+                    end
+
+      instance_definition = case instance_type
+                            when AST::Types::Name::Class, AST::Types::Name::Module
+                              type_name = checker.factory.type_name_1(instance_type.name)
+                              checker.factory.definition_builder.build_singleton(type_name)
+                            when AST::Types::Name::Instance
+                              type_name = checker.factory.type_name_1(instance_type.name)
+                              checker.factory.definition_builder.build_instance(type_name)
+                            end
+
+      module_definition = case module_type
+                          when AST::Types::Name::Class, AST::Types::Name::Module
+                            type_name = checker.factory.type_name_1(instance_type.name)
+                            checker.factory.definition_builder.build_singleton(type_name)
+                          else
+                            nil
+                          end
 
       module_context = TypeInference::Context::ModuleContext.new(
         instance_type: annots.instance_type || instance_type,
@@ -454,8 +477,8 @@ def for_sclass(node, type)
         current_namespace: current_namespace,
         const_env: self.module_context.const_env,
         class_name: self.module_context.class_name,
-        module_definition: nil,
-        instance_definition: self.module_context.module_definition
+        module_definition: module_definition,
+        instance_definition: instance_definition
       )
 
       type_env = TypeInference::TypeEnv.build(annotations: annots,
@@ -1177,8 +1200,10 @@ def synthesize(node, hint: nil)
 
             constructor.synthesize(node.children[1]) if node.children[1]
 
-            if constructor.module_context.instance_definition.type_name == self.module_context.module_definition.type_name
-              module_context.defined_module_methods.merge(constructor.module_context.defined_instance_methods)
+            if constructor.module_context.instance_definition && module_context.module_definition
+              if constructor.module_context.instance_definition.type_name == module_context.module_definition.type_name
+                module_context.defined_module_methods.merge(constructor.module_context.defined_instance_methods)
+              end
             end
 
             add_typing(node, type: AST::Builtin.nil_type)

test/type_construction_test.rb

@@ -4636,7 +4636,7 @@ class B
     end
   end
 
-  def test_singleton_class
+  def test_singleton_class_in_class_decl
     with_checker <<-RBS do |checker|
 class WithSingleton
   def self.open: [A] { () -> A } -> A
@@ -4672,4 +4672,87 @@ def open
       end
     end
   end
+
+  def test_singleton_class_in_class_decl_error
+    with_checker <<-RBS do |checker|
+class WithSingleton
+  def self.open: [A] { (instance) -> A } -> A
+end
+    RBS
+      source = parse_ruby(<<'EOF')
+class WithSingleton
+  class <<self
+    def open
+      yield 30
+    end
+  end
+end
+EOF
+
+      with_standard_construction(checker, source) do |construction, typing|
+        construction.synthesize(source.node)
+
+        assert_equal 1, typing.errors.size
+        assert_instance_of Steep::Errors::IncompatibleAssignment, typing.errors[0]
+      end
+    end
+  end
+
+  def test_singleton_class_for_object_success
+    with_checker <<-'RBS' do |checker|
+class WithSingleton
+  def open: [A] { () -> A } -> A
+end
+    RBS
+      source = parse_ruby(<<-'RUBY')
+class <<(WithSingleton.new)
+  def open
+    yield new()
+  end
+end
+      RUBY
+
+      with_standard_construction(checker, source) do |construction, typing|
+        type, _ = construction.synthesize(dig(source.node, 0))
+        sclass_constr = construction.for_sclass(dig(source.node), type)
+
+        module_context = sclass_constr.context.module_context
+
+        assert_equal parse_type("::WithSingleton"), module_context.instance_type
+        assert_equal parse_type("singleton(::WithSingleton)"), module_context.module_type
+        assert_nil module_context.class_name
+        assert_nil module_context.implement_name
+        assert_equal "::WithSingleton", module_context.module_definition.type_name.to_s
+        assert_equal "::WithSingleton", module_context.instance_definition.type_name.to_s
+
+        construction.synthesize(source.node)
+
+        assert_no_error typing
+      end
+    end
+  end
+
+  def test_singleton_class_for_object_type_check
+    with_checker <<-'RBS' do |checker|
+class WithSingleton
+  def open: [A] { () -> A } -> A
+end
+    RBS
+      source = parse_ruby(<<-'RUBY')
+class <<(WithSingleton.new)
+  def open(x)
+    x
+  end
+end
+      RUBY
+
+      with_standard_construction(checker, source) do |construction, typing|
+        construction.synthesize(source.node)
+
+        assert_equal 2, typing.errors.size
+        assert_instance_of Steep::Errors::MethodArityMismatch, typing.errors[0]
+        assert_instance_of Steep::Errors::FallbackAny, typing.errors[1]
+      end
+    end
+  end
 end