Add Class#hierarchy block support to allow formatting of classes

ManageIQ · Oct 13, 2023 · e612989 · e612989
1 parent 178de9e
commit e612989
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 11 deletions.
diff --git a/lib/more_core_extensions/core_ext/class/hierarchy.rb b/lib/more_core_extensions/core_ext/class/hierarchy.rb
@@ -22,15 +22,27 @@ def descendant_get(desc_name)
 
     # Returns a tree-like Hash structure of all descendants.
     #
+    #   If a block is passed, that block is used to format the keys of the hierarchy
+    #
     #   require 'socket'
     #   IO.hierarchy
     #   # => {BasicSocket=>
     #   #      {Socket=>{},
     #   #       IPSocket=>{TCPSocket=>{TCPServer=>{}}, UDPSocket=>{}},
     #   #       UNIXSocket=>{UNIXServer=>{}}},
     #   #     File=>{}}
-    def hierarchy
-      subclasses.each_with_object({}) { |k, h| h[k] = k.hierarchy }
+    #
+    #   IO.hierarchy(&:name)
+    #   # => {"BasicSocket"=>
+    #   #      {"Socket"=>{},
+    #   #       "IPSocket"=>{"TCPSocket"=>{"TCPServer"=>{}}, "UDPSocket"=>{}},
+    #   #       "UNIXSocket"=>{"UNIXServer"=>{}}},
+    #   #     "File"=>{}}
+    def hierarchy(&block)
+      subclasses.each_with_object({}) do |k, h|
+        key = block ? yield(k) : k
+        h[key] = k.hierarchy(&block)
+      end
     end
 
     # Returns an Array of all superclasses.

diff --git a/spec/core_ext/class/hierarchy_spec.rb b/spec/core_ext/class/hierarchy_spec.rb
@@ -27,15 +27,28 @@
     end
   end
 
-  it "#hierarchy" do
-    expect(IO.hierarchy).to eq(
-      BasicSocket => {
-        Socket     => {},
-        IPSocket   => {TCPSocket => {TCPServer => {}}, UDPSocket => {}},
-        UNIXSocket => {UNIXServer => {}}
-      },
-      File        => {}
-    )
+  describe "#hierarchy" do
+    it "by default returns a hash of descendants with keys being the classes" do
+      expect(IO.hierarchy).to eq(
+        BasicSocket => {
+          Socket     => {},
+          IPSocket   => {TCPSocket => {TCPServer => {}}, UDPSocket => {}},
+          UNIXSocket => {UNIXServer => {}}
+        },
+        File        => {}
+      )
+    end
+
+    it "with a block returns a hash of descendants with keys being the block evaluation" do
+      expect(IO.hierarchy(&:name)).to eq(
+        "BasicSocket" => {
+          "Socket"     => {},
+          "IPSocket"   => {"TCPSocket" => {"TCPServer" => {}}, "UDPSocket" => {}},
+          "UNIXSocket" => {"UNIXServer" => {}}
+        },
+        "File"        => {}
+      )
+    end
   end
 
   it "#lineage" do