From 8cebce6d1ac35f44ae0c7b7931991c613e3b1791 Mon Sep 17 00:00:00 2001
From: Robert Fletcher <lobatifricha@gmail.com>
Date: Wed, 12 Jun 2024 17:23:27 -0700
Subject: [PATCH] introduce BranchCollection class (#78)

This will encapsulate more logic for finding branches
---
 lib/baes.rb                    |  1 +
 lib/baes/actions/build_tree.rb |  6 +++--
 lib/baes/branch_collection.rb  | 40 ++++++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 lib/baes/branch_collection.rb

diff --git a/lib/baes.rb b/lib/baes.rb
index cc7b720..71fef54 100644
--- a/lib/baes.rb
+++ b/lib/baes.rb
@@ -10,5 +10,6 @@ class Baes::Error < StandardError; end
 
 require_relative "baes/actions"
 require_relative "baes/branch"
+require_relative "baes/branch_collection"
 require_relative "baes/git"
 require_relative "baes/version"
diff --git a/lib/baes/actions/build_tree.rb b/lib/baes/actions/build_tree.rb
index c80f88a..4827c24 100644
--- a/lib/baes/actions/build_tree.rb
+++ b/lib/baes/actions/build_tree.rb
@@ -36,9 +36,11 @@ def find_root_branch(branches)
     end
 
     def generate_branches
-      git.branch_names.map do |branch_name|
-        Baes::Branch.new(branch_name)
+      branches = Baes::BranchCollection.new
+      git.branch_names.each do |branch_name|
+        branches << Baes::Branch.new(branch_name)
       end
+      branches
     end
 
     def prune(branch)
diff --git a/lib/baes/branch_collection.rb b/lib/baes/branch_collection.rb
new file mode 100644
index 0000000..3d9a514
--- /dev/null
+++ b/lib/baes/branch_collection.rb
@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+# class to encapsulate a collection of branches
+class Baes::BranchCollection
+  attr_accessor :branches
+
+  def initialize
+    self.branches = []
+  end
+
+  # add a branch
+  def <<(branch)
+    branches << branch
+  end
+
+  # find a branch
+  def find(&block)
+    branches.find(&block)
+  end
+
+  # iterate over each branch with object
+  def each_with_object(obj, &block)
+    branches.each_with_object(obj, &block)
+  end
+
+  # iterate over each branch
+  def each(&block)
+    branches.each(&block)
+  end
+
+  # return the first branch
+  def first
+    branches.first
+  end
+
+  # return a branch by index
+  def [](index)
+    branches[index]
+  end
+end