From ff4eb64150ab716fbf047382b36e72325617063d Mon Sep 17 00:00:00 2001
From: Chad Wilken <chad.wilken@gmail.com>
Date: Wed, 1 Nov 2023 08:38:59 -0500
Subject: [PATCH] Implement enumerable for has_content module

---
 lib/tip_tap/has_content.rb    | 16 ++++++++++--
 spec/tip_tap/document_spec.rb | 47 +++++++++++++++++++++++++++++++----
 2 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/lib/tip_tap/has_content.rb b/lib/tip_tap/has_content.rb
index 7d56e36..48e45fd 100644
--- a/lib/tip_tap/has_content.rb
+++ b/lib/tip_tap/has_content.rb
@@ -4,6 +4,8 @@
 
 module TipTap
   module HasContent
+    include Enumerable
+
     attr_reader :attrs, :content
 
     def self.included(base)
@@ -20,13 +22,23 @@ def initialize(content = [], **attributes)
       yield self if block_given?
     end
 
-    def find_node(node_type)
-      content.find { |child| child.is_a?(node_type) }
+    def each
+      content.each { |child| yield child }
+    end
+
+    def find_node(type_class_or_name)
+      node_type = type_class_or_name.is_a?(String) ? TipTap.node_for(type_class_or_name) : type_class_or_name
+      find { |child| child.is_a?(node_type) }
     end
 
     def add_content(node)
       @content << node
     end
+    alias_method :<<, :add_content
+
+    def size
+      content.size
+    end
 
     module ClassMethods
       # Create a new instance from a TipTap JSON object.
diff --git a/spec/tip_tap/document_spec.rb b/spec/tip_tap/document_spec.rb
index d487706..5c1b690 100644
--- a/spec/tip_tap/document_spec.rb
+++ b/spec/tip_tap/document_spec.rb
@@ -80,13 +80,50 @@
     end
   end
 
+  describe "Enumerable" do
+    it "iterates over the content for #each" do
+      json = {
+        type: "doc",
+        content: [
+          {
+            type: "paragraph",
+            content: [
+              {
+                type: "text",
+                text: "Hello World!",
+                marks: [
+                  {type: "bold"},
+                  {type: "italic"}
+                ]
+              }
+            ]
+          }
+        ]
+      }
+      document = TipTap::Document.from_json(json)
+      expect(document.map(&:class)).to eq([TipTap::Nodes::Paragraph])
+    end
+  end
+
   describe "find_node" do
-    it "returns a node" do
-      document = TipTap::Document.from_json(json_contents)
-      node = document.find_node(TipTap::Nodes::Paragraph)
+    context "when passing a string" do
+      it "returns a node" do
+        document = TipTap::Document.from_json(json_contents)
+        node = document.find_node(TipTap::Nodes::Paragraph.type_name)
+
+        expect(node).to be_a(TipTap::Nodes::Paragraph)
+        expect(node.to_plain_text).to eq("Hello World!")
+      end
+    end
+
+    context "when passing a class" do
+      it "returns a node" do
+        document = TipTap::Document.from_json(json_contents)
+        node = document.find_node(TipTap::Nodes::Paragraph)
 
-      expect(node).to be_a(TipTap::Nodes::Paragraph)
-      expect(node.to_plain_text).to eq("Hello World!")
+        expect(node).to be_a(TipTap::Nodes::Paragraph)
+        expect(node.to_plain_text).to eq("Hello World!")
+      end
     end
   end