Spec: restrict the type returned by should_not be_nil and others

spec/std/spec/expectations_spec.cr

@@ -23,6 +23,20 @@ describe "expectations" do
     it { 100_000.should_not be_a(String) }
     it { 100_000.should be_a(Int32) }
     it { "Hello".should_not be_a(Int32) }
+
+    it "restricts type on should" do
+      x = 1 || 'a'
+      y = x.should be_a(Int32)
+      typeof(x).should eq(Int32 | Char)
+      typeof(y).should eq(Int32)
+    end
+
+    it "restricts type on should_not" do
+      x = 1 || 'a'
+      y = x.should_not be_a(Char)
+      typeof(x).should eq(Int32 | Char)
+      typeof(y).should eq(Int32)
+    end
   end
 
   describe "be_close" do
@@ -34,6 +48,13 @@ describe "expectations" do
     it { nil.should be_nil }
     it { "".should_not be_nil }
     it { 10.should_not be_nil }
+
+    it "restricts type on should_not" do
+      x = 1 || nil
+      y = x.should_not be_nil
+      typeof(x).should eq(Int32?)
+      typeof(y).should eq(Int32)
+    end
   end
 
   describe "be_falsey" do

src/spec/expectations.cr

@@ -417,6 +417,26 @@ module Spec
   end
 
   module ObjectExtensions
+    # Validates an expectation and fails the example if it does not match.
+    #
+    # This overload returns a value whose type is restricted to the expected type. For example:
+    #
+    # ```
+    # x = 1 || 'a'
+    # typeof(x) # => Int32 | Char
+    # x = x.should be_a(Int32)
+    # typeof(x) # => Int32
+    # ```
+    #
+    # See `Spec::Expecations` for available expectations.
+    def should(expectation : BeAExpectation(T), file = __FILE__, line = __LINE__) : T forall T
+      if expectation.match self
+        self.is_a?(T) ? self : (raise "Bug: expected #{self} to be a #{T}")
+      else
+        fail(expectation.failure_message(self), file, line)
+      end
+    end
+
     # Validates an expectation and fails the example if it does not match.
     #
     # See `Spec::Expecations` for available expectations.
@@ -426,6 +446,47 @@ module Spec
       end
     end
 
+    # Validates an expectation and fails the example if it matches.
+    #
+    # This overload returns a value whose type is restricted to exclude the given
+    # type in `should_not be_a`. For example:
+    #
+    # ```
+    # x = 1 || 'a'
+    # typeof(x) # => Int32 | Char
+    # x = x.should_not be_a(Char)
+    # typeof(x) # => Int32
+    # ```
+    #
+    # See `Spec::Expecations` for available expectations.
+    def should_not(expectation : BeAExpectation(T), file = __FILE__, line = __LINE__) forall T
+      if expectation.match self
+        fail(expectation.negative_failure_message(self), file, line)
+      else
+        self.is_a?(T) ? (raise "Bug: expected #{self} not to be a #{T}") : self
+      end
+    end
+
+    # Validates an expectation and fails the example if it matches.
+    #
+    # This overload returns a value whose type is restricted to be not `Nil`. For example:
+    #
+    # ```
+    # x = 1 || nil
+    # typeof(x) # => Int32 | Nil
+    # x = x.should_not be_nil
+    # typeof(x) # => Int32
+    # ```
+    #
+    # See `Spec::Expecations` for available expectations.
+    def should_not(expectation : BeNilExpectation, file = __FILE__, line = __LINE__)
+      if expectation.match self
+        fail(expectation.negative_failure_message(self), file, line)
+      else
+        self.not_nil!
+      end
+    end
+
     # Validates an expectation and fails the example if it matches.
     #
     # See `Spec::Expecations` for available expectations.