From e7b2ff9392e1c7e4da29678014ca2244f63b97d0 Mon Sep 17 00:00:00 2001
From: Sam Bostock <sam.bostock@shopify.com>
Date: Tue, 16 Aug 2022 18:57:14 -0400
Subject: [PATCH] Add `AllowOnSelfClass` option to `Style/CaseEquality`

`Style/CaseEquality`'s `AllowOnConstant` option allows the use of constants, to
allow for case equality with classes and modules. `self.class` is a special
case also referring to a class, and may be required in cases where the other
object might be a `BasicObject`.
---
 .../new_add_allowonselfclass_option_to.md     |  1 +
 config/default.yml                            |  9 ++++
 lib/rubocop/cop/style/case_equality.rb        | 50 +++++++++++++++----
 spec/rubocop/cop/style/case_equality_spec.rb  | 38 ++++++++++++++
 4 files changed, 88 insertions(+), 10 deletions(-)
 create mode 100644 changelog/new_add_allowonselfclass_option_to.md

diff --git a/changelog/new_add_allowonselfclass_option_to.md b/changelog/new_add_allowonselfclass_option_to.md
new file mode 100644
index 00000000000..72e52260dd2
--- /dev/null
+++ b/changelog/new_add_allowonselfclass_option_to.md
@@ -0,0 +1 @@
+* [#10931](https://github.com/rubocop/rubocop/pull/10931): Add `AllowOnSelfClass` option to `Style/CaseEquality`. ([@sambostock][])
diff --git a/config/default.yml b/config/default.yml
index e4e1778b521..412166f1d26 100644
--- a/config/default.yml
+++ b/config/default.yml
@@ -3186,6 +3186,15 @@ Style/CaseEquality:
   #   # good
   #   String === "string"
   AllowOnConstant: false
+  # If `AllowOnSelfClass` option is enabled, the cop will ignore violations when the receiver of
+  # the case equality operator is `self.class`.
+  #
+  #   # bad
+  #   some_class === object
+  #
+  #   # good
+  #   self.class === object
+  AllowOnSelfClass: false
 
 Style/CaseLikeIf:
   Description: 'Identifies places where `if-elsif` constructions can be replaced with `case-when`.'
diff --git a/lib/rubocop/cop/style/case_equality.rb b/lib/rubocop/cop/style/case_equality.rb
index a5f6cd45231..3e841c56cc9 100644
--- a/lib/rubocop/cop/style/case_equality.rb
+++ b/lib/rubocop/cop/style/case_equality.rb
@@ -7,6 +7,9 @@ module Style
       #
       # If `AllowOnConstant` option is enabled, the cop will ignore violations when the receiver of
       # the case equality operator is a constant.
+
+      # If `AllowOnSelfClass` option is enabled, the cop will ignore violations when the receiver of
+      # the case equality operator is `self.class`. Note intermediate variables are not accepted.
       #
       # @example
       #   # bad
@@ -26,6 +29,14 @@ module Style
       #   # good
       #   Array === something
       #
+      # @example AllowOnSelfClass: false (default)
+      #   # bad
+      #   self.class === something
+      #
+      # @example AllowOnSelfClass: true
+      #   # good
+      #   self.class === something
+      #
       class CaseEquality < Base
         extend AutoCorrector
 
@@ -33,7 +44,10 @@ class CaseEquality < Base
         RESTRICT_ON_SEND = %i[===].freeze
 
         # @!method case_equality?(node)
-        def_node_matcher :case_equality?, '(send $#const? :=== $_)'
+        def_node_matcher :case_equality?, '(send $#offending_receiver? :=== $_)'
+
+        # @!method self_class?(node)
+        def_node_matcher :self_class?, '(send (self) :class)'
 
         def on_send(node)
           case_equality?(node) do |lhs, rhs|
@@ -48,12 +62,11 @@ def on_send(node)
 
         private
 
-        def const?(node)
-          if cop_config.fetch('AllowOnConstant', false)
-            !node&.const_type?
-          else
-            true
-          end
+        def offending_receiver?(node)
+          return false if node&.const_type? && cop_config.fetch('AllowOnConstant', false)
+          return false if self_class?(node) && cop_config.fetch('AllowOnSelfClass', false)
+
+          true
         end
 
         def replacement(lhs, rhs)
@@ -66,12 +79,29 @@ def replacement(lhs, rhs)
             #
             # So here is noop.
           when :begin
-            child = lhs.children.first
-            "#{lhs.source}.include?(#{rhs.source})" if child&.range_type?
+            begin_replacement(lhs, rhs)
           when :const
-            "#{rhs.source}.is_a?(#{lhs.source})"
+            const_replacement(lhs, rhs)
+          when :send
+            send_replacement(lhs, rhs)
           end
         end
+
+        def begin_replacement(lhs, rhs)
+          return unless lhs.children.first&.range_type?
+
+          "#{lhs.source}.include?(#{rhs.source})"
+        end
+
+        def const_replacement(lhs, rhs)
+          "#{rhs.source}.is_a?(#{lhs.source})"
+        end
+
+        def send_replacement(lhs, rhs)
+          return unless self_class?(lhs)
+
+          "#{rhs.source}.is_a?(#{lhs.source})"
+        end
       end
     end
   end
diff --git a/spec/rubocop/cop/style/case_equality_spec.rb b/spec/rubocop/cop/style/case_equality_spec.rb
index 9e66acd623c..39960f809b7 100644
--- a/spec/rubocop/cop/style/case_equality_spec.rb
+++ b/spec/rubocop/cop/style/case_equality_spec.rb
@@ -73,4 +73,42 @@
 
     include_examples 'offenses'
   end
+
+  context 'when AllowOnSelfClass is false' do
+    let(:cop_config) { { 'AllowOnSelfClass' => false } }
+
+    it 'registers an offense and corrects for === when the receiver is self.class' do
+      expect_offense(<<~RUBY)
+        self.class === var
+                   ^^^ Avoid the use of the case equality operator `===`.
+      RUBY
+
+      expect_correction(<<~RUBY)
+        var.is_a?(self.class)
+      RUBY
+    end
+
+    include_examples 'offenses'
+  end
+
+  context 'when AllowOnSelfClass is true' do
+    let(:cop_config) { { 'AllowOnSelfClass' => true } }
+
+    it 'does not register an offense for === when the receiver is self.class' do
+      expect_no_offenses(<<~RUBY)
+        self.class === var
+      RUBY
+    end
+
+    it 'registers an offense and corrects for === when the receiver is self.klass' do
+      expect_offense(<<~RUBY)
+        self.klass === var
+                   ^^^ Avoid the use of the case equality operator `===`.
+      RUBY
+
+      expect_no_corrections
+    end
+
+    include_examples 'offenses'
+  end
 end