Handle 'self' and some other variants in method return types

CHANGELOG.md

@@ -6,7 +6,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 ### Fixed
-- Fix macro logic to support free variables on stubbed methods.
+- Fix macro logic to support free variables, 'self', and variants on stubbed methods. [#48](https://github.com/icy-arctic-fox/spectator/issues/48)
 
 ### Changed
 - Simplify string representation of mock-related types.

spec/issues/github_issue_48_spec.cr

@@ -13,6 +13,22 @@ Spectator.describe "GitHub Issue #48" do
     def make_nilable(thing : T) : T? forall T
       thing.as(T?)
     end
+
+    def itself : self
+      self
+    end
+
+    def itself? : self?
+      self.as(self?)
+    end
+
+    def generic(thing : T) : Array(T) forall T
+      Array.new(100) { thing }
+    end
+
+    def union : Int32 | String
+      42.as(Int32 | String)
+    end
   end
 
   mock Test, make_nilable: nil
@@ -40,7 +56,43 @@ Spectator.describe "GitHub Issue #48" do
   end
 
   it "handles nilable free variables" do
-    fake = mock(Test)
     expect(fake.make_nilable("foo")).to be_nil
   end
+
+  it "handles 'self' return type" do
+    not_self = mock(Test)
+    allow(fake).to receive(:itself).and_return(not_self)
+    expect(fake.itself).to be(not_self)
+  end
+
+  it "raises on type cast error with 'self' return type" do
+    allow(fake).to receive(:itself).and_return(42)
+    expect { fake.itself }.to raise_error(TypeCastError, /#{class_mock(Test)}/)
+  end
+
+  it "handles nilable 'self' return type" do
+    not_self = mock(Test)
+    allow(fake).to receive(:itself?).and_return(not_self)
+    expect(fake.itself?).to be(not_self)
+  end
+
+  it "handles generic return type" do
+    allow(fake).to receive(:generic).and_return([42])
+    expect(fake.generic(42)).to eq([42])
+  end
+
+  it "raises on type cast error with generic return type" do
+    allow(fake).to receive(:generic).and_return("test")
+    expect { fake.generic(42) }.to raise_error(TypeCastError, /Array\(Int32\)/)
+  end
+
+  it "handles union return types" do
+    allow(fake).to receive(:union).and_return("test")
+    expect(fake.union).to eq("test")
+  end
+
+  it "raises on type cast error with union return type" do
+    allow(fake).to receive(:union).and_return(:test)
+    expect { fake.union }.to raise_error(TypeCastError, /Symbol/)
+  end
 end

src/spectator/mocks/stubbable.cr

@@ -158,12 +158,24 @@ module Spectator
           # Cast the stub or return value to the expected type.
           # This is necessary to match the expected return type of the original method.
           _spectator_cast_stub_value(%stub, %call, typeof({{original}}),
-          {{ if method.return_type && method.return_type.resolve? == NoReturn
-               :no_return
-             elsif method.return_type &&
-                   ((resolved = method.return_type.resolve?).is_a?(TypeNode) && resolved <= Nil) ||
-                   (method.return_type.is_a?(Union) && method.return_type.types.map(&.resolve?).includes?(Nil))
-               :nil
+          {{ if rt = method.return_type
+               if rt.is_a?(Path) && (resolved = rt.resolve?).is_a?(TypeNode) && resolved <= NoReturn
+                 :no_return
+               else
+                 # Process as an enumerable type to reduce code repetition.
+                 rt = rt.is_a?(Union) ? rt.types : [rt]
+                 # Check if any types are nilable.
+                 nilable = rt.any? do |t|
+                   # These are all macro types that have the `resolve?` method.
+                   (t.is_a?(TypeNode) || t.is_a?(Path) || t.is_a?(Generic) || t.is_a?(MetaClass)) &&
+                     (resolved = t.resolve?).is_a?(TypeNode) && resolved <= Nil
+                 end
+                 if nilable
+                   :nil
+                 else
+                   :raise
+                 end
+               end
              else
                :raise
              end }})
@@ -261,16 +273,25 @@ module Spectator
         if %stub = _spectator_find_stub(%call)
           # Cast the stub or return value to the expected type.
           # This is necessary to match the expected return type of the original method.
-          {% if method.return_type %}
+          {% if rt = method.return_type %}
             # Return type restriction takes priority since it can be a superset of the original implementation.
             _spectator_cast_stub_value(%stub, %call, {{method.return_type}},
-              {{ if method.return_type.resolve? == NoReturn
+              {{ if rt.is_a?(Path) && (resolved = rt.resolve?).is_a?(TypeNode) && resolved <= NoReturn
                    :no_return
-                 elsif (method.return_type.resolve?.is_a?(TypeNode) && method.return_type.resolve <= Nil) ||
-                       (method.return_type.is_a?(Union) && method.return_type.types.map(&.resolve?).includes?(Nil))
-                   :nil
                  else
-                   :raise
+                   # Process as an enumerable type to reduce code repetition.
+                   rt = rt.is_a?(Union) ? rt.types : [rt]
+                   # Check if any types are nilable.
+                   nilable = rt.any? do |t|
+                     # These are all macro types that have the `resolve?` method.
+                     (t.is_a?(TypeNode) || t.is_a?(Path) || t.is_a?(Generic) || t.is_a?(MetaClass)) &&
+                       (resolved = t.resolve?).is_a?(TypeNode) && resolved <= Nil
+                   end
+                   if nilable
+                     :nil
+                   else
+                     :raise
+                   end
                  end }})
           {% elsif !method.abstract? %}
             # The method isn't abstract, infer the type it returns without calling it.