lib/druid/elements/table_row.rb in druid-ts-1.2.4 vs lib/druid/elements/table_row.rb in druid-ts-1.2.5

@@ -6,37 +6,41 @@
       #
       # Return the Druid::Elements::TableCell for the index provided. Index
       # is zero based. If the index provided is a String then it
       # will be matched with the text from the columns in the first row.
       #
-      def [](idx)
-        idx = find_index_by_title(idx) if idx.kind_of?(String)
-        return nil unless idx && columns >= idx + 1
-        Druid::Elements::TableCell.new(element[idx])
+      def [](what)
+        idx = find_index(what)
+        idx && cell_items[idx]
       end
+
       #
-      # Returns the number of rows in the table.
+      # Returns the number of colums in the table
       #
       def columns
-        element.cells.size
+        cell_items.size
       end
 
-      def each
-        for index in 1..self.columns do
-          yield self[index-1]
-        end
+      #
+      # iterator that yields with a Druid::Elements::TableCell
+      #
+      def each(&block)
+        cell_items.each(&block)
       end
 
       private
 
-      def find_index_by_title(title)
-        table = element.parent
-        table = table.parent if table.tag_name == 'tbody'
-        if table.instance_of? Watir::HTMLElement
-          table = table.to_subtype
+      def cell_items
+        @cell_items ||= element.cells.map do |obj|
+          Druid::Elements::TableCell.new(obj)
         end
-        first_row = table[0]
-        first_row.cells.find_index { |column| column.text.include? title}
+      end
+
+      def find_index(what)
+        return what if what.is_a? Integer
+        parent(tag_name: 'table').headers.find_index do |header|
+          header.text.include? what
+        end
       end
 
     end
 
     Druid::Elements.tag_to_class[:tr] = Druid::Elements::TableRow