From 37ef3865afc85a720dfa17a775543d2bf4c8ae8a Mon Sep 17 00:00:00 2001
From: Teppei Shintani <teppest1215@gmail.com>
Date: Sun, 7 Jan 2024 18:38:19 +0900
Subject: [PATCH] Support nested array and nested hash

---
 lib/packwerk_yard/parser.rb            | 14 ++++++++++----
 test/fixtures/yard/nested_array.rb     |  9 +++++++++
 test/fixtures/yard/nested_hash.rb      |  9 +++++++++
 test/unit/packwerk_yard/parser_test.rb | 21 +++++++++++++++++++++
 4 files changed, 49 insertions(+), 4 deletions(-)
 create mode 100644 test/fixtures/yard/nested_array.rb
 create mode 100644 test/fixtures/yard/nested_hash.rb

diff --git a/lib/packwerk_yard/parser.rb b/lib/packwerk_yard/parser.rb
index ec0451f..87cc8fd 100644
--- a/lib/packwerk_yard/parser.rb
+++ b/lib/packwerk_yard/parser.rb
@@ -7,9 +7,13 @@ class Parser
     include Packwerk::FileParser
 
     # Array Syntax e.g. Array<String>
-    ARRAY_REGEXP = T.let(/Array<(.+)>/.freeze, Regexp)
+    ARRAY_REGEXP = T.let(/\AArray<(.+)>/.freeze, Regexp)
     private_constant :ARRAY_REGEXP
 
+    # Hash Syntax e.g. Hash<String, String>
+    HASH_REGEXP = T.let(/\AHash<([^,]*),\s?(.*)>/.freeze, Regexp)
+    private_constant :HASH_REGEXP
+
     sig { params(ruby_parser: T.nilable(Packwerk::Parsers::Ruby)).void }
     def initialize(ruby_parser: Packwerk::Parsers::Ruby.new)
       @ruby_parser = ruby_parser
@@ -23,7 +27,7 @@ def call(io:, file_path: "<unknown>")
       types = extract_from_yard_to_types(source_code)
 
       to_ruby_ast(
-        types.map { |type| to_evaluable_type(type) }
+        types.map { |type| to_evaluable_type(type) }.flatten
              .reject { |type| to_constant(type).nil? }
              .inspect.delete('"'),
         file_path,
@@ -56,9 +60,11 @@ def extract_from_yard_to_types(source_code)
       types.uniq
     end
 
-    sig { params(type: String).returns(String) }
+    sig { params(type: String).returns(T::Array[String]) }
     def to_evaluable_type(type)
-      ARRAY_REGEXP.match(type).to_a[1] || 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.empty? ? [type] : matched_types.map {|t | to_evaluable_type(t) }.flatten
     end
 
     sig { params(name: T.any(Symbol, String)).returns(T.untyped) }
diff --git a/test/fixtures/yard/nested_array.rb b/test/fixtures/yard/nested_array.rb
new file mode 100644
index 0000000..9b5c80f
--- /dev/null
+++ b/test/fixtures/yard/nested_array.rb
@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class NestedArray
+  # @param [Array<Array<String>>] array
+  # @return [Array<Array<String>>]
+  def nested_array_method(array)
+    array
+  end
+end
diff --git a/test/fixtures/yard/nested_hash.rb b/test/fixtures/yard/nested_hash.rb
new file mode 100644
index 0000000..4718e39
--- /dev/null
+++ b/test/fixtures/yard/nested_hash.rb
@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class NestedHash
+  # @param [Hash<Integer, Hash<Integer, String>>] hash
+  # @return [Hash<Integer, Hash<Integer, String>>]
+  def nested_hash(hash)
+    hash
+  end
+end
diff --git a/test/unit/packwerk_yard/parser_test.rb b/test/unit/packwerk_yard/parser_test.rb
index 0842e02..3ed42ba 100644
--- a/test/unit/packwerk_yard/parser_test.rb
+++ b/test/unit/packwerk_yard/parser_test.rb
@@ -24,5 +24,26 @@ def test_call_return_node_with_not_exists_class_args_file
       assert(node.type, :array)
       assert(node.children.size, 0)
     end
+
+    def test_call_return_node_with_nested_array_file
+      parser = ::PackwerkYard::Parser.new
+      io = StringIO.new(File.read("test/fixtures/yard/nested_array.rb"))
+      node = parser.call(io: io, file_path: "test/fixtures/yard/nested_array.rb")
+
+      assert_instance_of(::Parser::AST::Node, node)
+      assert(node.type, :array)
+      assert(node.children[0].children[1], :String)
+    end
+
+    def test_call_return_node_with_nested_hash_file
+      parser = ::PackwerkYard::Parser.new
+      io = StringIO.new(File.read("test/fixtures/yard/nested_hash.rb"))
+      node = parser.call(io: io, file_path: "test/fixtures/yard/nested_hash.rb")
+
+      assert_instance_of(::Parser::AST::Node, node)
+      assert(node.type, :array)
+      assert(node.children[0].children[1], :String)
+      assert(node.children[1].children[1], :Integer)
+    end
   end
 end