lib/watir-webdriver/collections/element_collection.rb in watir-webdriver-0.0.6 vs lib/watir-webdriver/collections/element_collection.rb in watir-webdriver-0.0.7

@@ -1,16 +1,17 @@
 # encoding: utf-8
 module Watir
   class ElementCollection
     include Enumerable
 
-    def initialize(parent, element_class)
-      @parent, @element_class = parent, element_class
+    def initialize(parent, selector)
+      @parent   = parent
+      @selector = selector
     end
 
     #
-    # @yieldparam [Watir::BaseElement] element Iterate through the elements in this collection.
+    # @yieldparam [Watir::Element] element Iterate through the elements in this collection.
     #
 
     def each(&blk)
       to_a.each(&blk)
     end
@@ -30,56 +31,61 @@
     #
     # Also note that because of Watir's lazy loading, this will return an Element
     # instance even if the index is out of bounds.
     #
     # @param [Fixnum] n Index of wanted element, 0-indexed
-    # @return [Watir::BaseElement] Returns an instance of a Watir::BaseElement subclass
+    # @return [Watir::Element] Returns an instance of a Watir::Element subclass
     #
 
 
     def [](idx)
-      to_a[idx] || @element_class.new(@parent, :index, idx)
+      to_a[idx] || element_class.new(@parent, :index => idx)
     end
 
     #
     # First element of this collection
     #
-    # @return [Watir::BaseElement] Returns an instance of a Watir::BaseElement subclass
+    # @return [Watir::Element] Returns an instance of a Watir::Element subclass
     #
 
     def first
-      to_a[0] || @element_class.new(@parent, :index, 0)
+      to_a[0] || element_class.new(@parent, :index => 0)
     end
 
     #
     # Last element of the collection
     #
-    # @return [Watir::BaseElement] Returns an instance of a Watir::BaseElement subclass
+    # @return [Watir::Element] Returns an instance of a Watir::Element subclass
     #
 
     def last
-      to_a[-1] || @element_class.new(@parent, :index, -1)
+      to_a[-1] || element_class.new(@parent, :index => -1)
     end
 
     #
     # This collection as an Array
     #
-    # @return [Array<Watir::BaseElement>]
+    # @return [Array<Watir::Element>]
     #
 
     def to_a
-      # TODO: optimize - lazily @element_class instance
-      @to_a ||= elements.map { |e| @element_class.new(@parent, :element, e) }
+      # TODO: optimize - lazy element_class instance?
+      @to_a ||= elements.map { |e| element_class.new(@parent, :element => e) }
     end
 
     private
 
     def elements
-      @elements ||= ElementLocator.new(
+      @elements ||= locator_class.new(
         @parent.wd,
-        @element_class.default_selector,
-        @element_class.attribute_list
+        @selector,
+        element_class.attribute_list
       ).locate_all
+    end
+
+    # overridable by subclasses
+    def locator_class
+      ElementLocator
     end
 
   end # ElementCollection
 end # Watir