document.rb in rubyfocus-0.4.0

- old
+ new

@@ -51,11 +51,11 @@
 	def self.from_local
 		new(Rubyfocus::LocalFetcher.new)
 	end
 
 	# Initialize with a URL, for remote fetching.
-	# Not implemented yet
+	# Not implemented yet TODO implement
 	def self.from_url(url)
 		raise RuntimeError, "Rubyfocus::Document.from_url not yet implemented."
 		# new(Rubyfocus::RemoteFetcher.new(url))
 	end
 
@@ -94,28 +94,41 @@
 			Rubyfocus::Setting => @settings
 		}[obj.class]
 	end
 	private :ivar_for
 
-	# Add an element. Element should be a Project, Task, Context, Folder, or Setting
-	# We assume whoever does this will set document appropriately on the element
-	def add_element(e)
+	# Add an element. Element should be a Project, Task, Context, Folder, or Setting.
+	# If overwrite set to false and ID already occurs in the document, throw an error.
+	# If ID is nil, throw an error.
+	def add_element(e, overwrite:false)
+		# Error check
+		raise(Rubyfocus::DocumentElementException, "Adding element to document, but it has no ID.") if e.id.nil?
+		raise(Rubyfocus::DocumentElementException, "Adding element to document, but element with this ID already exists.") if !overwrite && has_id?(e.id)
+
+		# Otherwise, full steam ahead
+		e.document = self
+
+		if (dupe_element = self[e.id]) && overwrite
+			remove_element(dupe_element)
+		end
+
+		# Add to the correct array
 		dest = ivar_for(e)
 		if dest
 			dest << e
 		else
 			raise ArgumentError, "You passed a #{e.class} to Document#add_element - I don't know what to do with this."
 		end
 	end
 
 	# Remove an element from the document.
-	# We assume whoever does this is smart enough to also set the element's #document value
-	# to nil
 	def remove_element(e)
 		e = self[e] if e.is_a?(String)
 		return if e.nil?
 
+		e.document = nil
+
 		dest = ivar_for(e)
 		if dest
 			dest.delete(e)
 		else
 			raise ArgumentError, "You passed a #{e.class} to Document#remove_element - I don't know what to do with this."	
@@ -129,18 +142,28 @@
 	end
 
 	# For Searchable include
 	alias_method :array, :elements
 
+
 	#-------------------------------------------------------------------------------
 	# Find elements from id
 	def [] search_id
 		self.elements.find{ |elem| elem.id == search_id }
 	end
 
+	# Check if the document has an element of a given ID
+	def has_id?(id)
+		self.elements.any?{ |e| e.id == id }
+	end
+
 	#---------------------------------------
 	# YAML export
 
 	def save(file)
 		File.open(file, "w"){ |io| io.puts YAML::dump(self) }
 	end
-end
+end
+
+#-------------------------------------------------------------------------------
+# Exceptions
+class Rubyfocus::DocumentElementException < Exception; end
\ No newline at end of file