From ca898ea102db27bfb6f79b05e9817244e94a6ff5 Mon Sep 17 00:00:00 2001
From: Teppei Shintani <teppest1215@gmail.com>
Date: Sun, 7 Jan 2024 23:32:08 +0900
Subject: [PATCH] Support hash specific format sush as Hash{String => Integer}

---
 lib/packwerk_yard/constantize_type.rb           |  7 ++++++-
 .../unit/packwerk_yard/constantize_type_test.rb |  8 ++++++++
 test/unit/packwerk_yard/yard_handler_test.rb    | 17 +++++++++++++++++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/lib/packwerk_yard/constantize_type.rb b/lib/packwerk_yard/constantize_type.rb
index ed5b014..6b553a0 100644
--- a/lib/packwerk_yard/constantize_type.rb
+++ b/lib/packwerk_yard/constantize_type.rb
@@ -10,9 +10,13 @@ class ConstantizeType
     private_constant :ARRAY_REGEXP
 
     # Hash Syntax e.g. Hash<String, String>
-    HASH_REGEXP = T.let(/\AHash<([^,]*),\s?(.*)>/.freeze, Regexp)
+    HASH_REGEXP = T.let(/\AHash<([^,\s]+)\s*,\s*(.+)>/.freeze, Regexp)
     private_constant :HASH_REGEXP
 
+    # Hash Syntax e.g. Hash{String => String}
+    HASH_SPECIFIC_REGEXP = T.let(/\AHash\{([^,\s]+)\s*=>\s*(.+)}/.freeze, Regexp)
+    private_constant :HASH_SPECIFIC_REGEXP
+
     sig { params(yard_type: String).void }
     def initialize(yard_type)
       @yard_type = yard_type
@@ -37,6 +41,7 @@ def constantize?(name)
     def split_type(type)
       matched_types = Array(ARRAY_REGEXP.match(type).to_a[1])
       matched_types = Array(HASH_REGEXP.match(type).to_a[1..2]) if matched_types.empty?
+      matched_types = Array(HASH_SPECIFIC_REGEXP.match(type).to_a[1..2]) if matched_types.empty?
       matched_types.empty? ? [type] : matched_types.map { |t| split_type(t) }.flatten
     end
   end
diff --git a/test/unit/packwerk_yard/constantize_type_test.rb b/test/unit/packwerk_yard/constantize_type_test.rb
index 930d7ef..8cc2ba8 100644
--- a/test/unit/packwerk_yard/constantize_type_test.rb
+++ b/test/unit/packwerk_yard/constantize_type_test.rb
@@ -36,6 +36,14 @@ def test_types_returns_constantize_types_with_nested_hash
       assert_equal(["Integer", "String"], types)
     end
 
+    def test_types_returns_constantize_types_with_hash_specific_format
+      types = PackwerkYard::ConstantizeType.new("Hash{Integer => String}").types
+      assert_equal(["Integer", "String"], types)
+
+      types = PackwerkYard::ConstantizeType.new("Hash{Integer=>String}").types
+      assert_equal(["Integer", "String"], types)
+    end
+
     def test_types_returns_constantize_types_with_not_existing_class
       types = PackwerkYard::ConstantizeType.new("NotExistsClass").types
 
diff --git a/test/unit/packwerk_yard/yard_handler_test.rb b/test/unit/packwerk_yard/yard_handler_test.rb
index 0e161a6..0b138ec 100644
--- a/test/unit/packwerk_yard/yard_handler_test.rb
+++ b/test/unit/packwerk_yard/yard_handler_test.rb
@@ -68,5 +68,22 @@ def bar(foo)
       assert_equal(["String", "Integer"], handler.param_types)
       assert_equal(["String", "Integer"], handler.return_types)
     end
+
+    def test_from_source_with_hash_specific_format
+      code = <<~RUBY
+        class Foo
+          # @param [Hash{Integer => String}] foo
+          # @return [Hash{Integer=>String}]
+          def bar(foo)
+            foo
+          end
+        end
+      RUBY
+
+      handler = PackwerkYard::YardHandler.from_source(code)
+
+      assert_equal(["Hash{Integer => String}"], handler.param_types)
+      assert_equal(["Hash{Integer=>String}"], handler.return_types)
+    end
   end
 end