lib/rgraph/graph.rb in rgraph-0.0.3 vs lib/rgraph/graph.rb in rgraph-0.0.4

@@ -7,30 +7,39 @@
   attr_accessor :nodes, :links
 
   def initialize(csv)
     @nodes = []
     @links = []
-    unless not csv[-4..-1] == ".csv"
-      CSV.foreach(csv, headers: true) do |row|
-        unless source = @nodes.select{|n| n.id == row['source']}.first
-          source = Node.new(id: row['source'])
-          @nodes << source
-        end
-        unless target = @nodes.select{|n| n.id == row['target']}.first
-          target = Node.new(id: row['target'])
-          @nodes << target
-        end
-        @links << Link.new(source: source, target: target, weight: row['weight'], year: row['year'])
+    raise Exception.new("the file must be a .csv") unless File.extname(csv) == ".csv"
+
+    CSV.foreach(csv, headers: true) do |row|
+      #last because CSV#delete returns [column,value]
+      source_id = row.delete('source').last
+      target_id = row.delete('target').last
+
+      unless source = get_node_by_id(source_id)
+        source = Node.new(id: source_id)
+        @nodes << source
       end
-    else
-      raise Exception.new("the file must be a .csv")
+      unless target = get_node_by_id(target_id)
+        target = Node.new(id: target_id)
+        @nodes << target
+      end
+
+      @links << Link.new(source: source, target: target, weight: row['weight'], year: row['year'])
     end
-end
+  end
 
   def each_node(&block)
     @nodes.each(&block)
   end
 
   def each_link(&block)
     @links.each(&block)
+  end
+
+  private
+
+  def get_node_by_id(node_id)
+    @nodes.select{|n| n.id == node_id}.first
   end
 end