Improve color validation on shapes

2 files changed, 22 insertions, 5 deletions

diff --git a/lib/ruby2d/quad.rb b/lib/ruby2d/quad.rb
index eca123d..f17aa5c 100644
--- a/lib/ruby2d/quad.rb
+++ b/lib/ruby2d/quad.rb
@@ -47,16 +47,22 @@ module Ruby2D
         @c3 = Color.new(c)
         @c4 = Color.new(c)
         
+      elsif c.class == Array && c.length < 4
+        raise Error, "Quads require 4 colors, one for each vertex. Only " <<
+                     "#{c.length} were given."
+        
       # If a valid array of colors, assign them to each vertex, respectively
       elsif c.all? { |el| Color.is_valid? el }
         @c1 = Color.new(c[0])
         @c2 = Color.new(c[1])
         @c3 = Color.new(c[2])
         @c4 = Color.new(c[3])
+        
       else
         raise Error, "Not a valid color for #{self.class}"
       end
       
     end
+    
   end
 end
diff --git a/lib/ruby2d/triangle.rb b/lib/ruby2d/triangle.rb
index ad13688..ac7d67a 100644
--- a/lib/ruby2d/triangle.rb
+++ b/lib/ruby2d/triangle.rb
@@ -35,16 +35,27 @@ module Ruby2D
     private
     
     def update_color(c)
-      # If a 2D array
-      if c.class == Array && c.all? { |el| el.class == Array }
+      
+      # If a valid color, use it for each vertex
+      if Color.is_valid? c
+        @c1 = Color.new(c)
+        @c2 = Color.new(c)
+        @c3 = Color.new(c)
+        
+      elsif c.class == Array && c.length < 3
+        raise Error, "Triangles require 3 colors, one for each vertex. Only " <<
+                     "#{c.length} were given."
+        
+      # If a valid array of colors, assign them to each vertex, respectively
+      elsif c.all? { |el| Color.is_valid? el }
         @c1 = Color.new(c[0])
         @c2 = Color.new(c[1])
         @c3 = Color.new(c[2])
+        
       else
-        @c1 = Color.new(c)
-        @c2 = Color.new(c)
-        @c3 = Color.new(c)
+        raise Error, "Not a valid color for #{self.class}"
       end
+      
     end
     
   end