[pygtk] XML and gtk.TreeStores

[Seth Mahoney]

Hey all, for a project I'm working on, I subclassed gtk.TreeStore so
that it can automatically load data from an xml file (using lxml), and
then subclassed gtk.TreeView so that it automatically sets up
gtk.TreeViewColumns and gtk.CellRenderers based on the data that's
passed to it.  If it comes in handy for anyone, feel free to use it, and
if anyone fixes any bugs or finds any speed fixes (or nicer ways to get
stuff done - some of it looks pretty hackish to me), let me know.

--Seth


#!/usr/bin/env python
import pygtk, gtk
from lxml import etree
from lxml.etree import Element, SubElement

#This class will automatically load xml data into a treestore.
#It expects that the arguments passed to its constructor begin
#with a xml_file, which can be a filename or an etree Element,
#and a bool indicating whether or not you want the
#class to update the xml every time the tree is updated,
#and then include pairs of xml tag names and types.
#So, assuming you started with an xml tree that looked like:
#<tree>
#	<item>
#		<data1>Voltaire</data1>
#		<data2>False</data2>
#	</item>
#</tree>
#where data1 is a string which holds a name and data2 a bool
#that says whether or not that person lives, you would call the
#constructor like so:
#xml_tree = xml_tree_store("xmlfile.xml", True, "Name", str, "Alive", bool)
#and the treestore would be loaded with two columns, "Name" and "Alive",
#which hold a name and a bool, respectively, and each row in the treestore
#would correspond to an xml item, as designated by the <item> tag in the
#tree above (this tag doesn't need to be called "item", and the name is
#determined automatically).  One caveat: at this time, only str, bool, and
#int types are supported.  Any other type will be handled as a str.  Further,
#bool types should be represented in the xml file as "true" (without the quotes)
#and "false", using whatever capitalization scheme seems appropriate.  Finally,
#all data should be stored between the tags (<tag>Here's some data!</tag>) and
#not as attributes (<tag data="Some data!"/>.  More complicated tree structures
#are supported, such as:
#<tree>
#	<item>
#		<data1/>
#		<data2/>
#		<item>
#			<data1/>
#			<data2/>
#		</item>
#	</item>
#</tree>
#and will be handled correctly, with the second item stored as the first
#item's child in the treestore.  The class does expect, however, that
#the xml structure will be roughly as follows:
#<tree>
#	<item/>
#	<item/>
#	<item/>
#</tree>
#with only items (or whatever you decide to call them) as subelements of
#the tree.  If subelements with different tags are stored as children of
#the root element (which doesn't need to be called "tree"), only those
#with the same tag as the first subelement will be read.
#Finally, an extra, hidden column will always be added to the end of the
#treestore, which will hold an xml element and which can be used, if
#necessary, to traverse the xml tree, although generally the standard
#gtk functions should be used instead (any functions which modify the
#TreeStore will also modify the xml tree unless you set update_xml to false.

class xml_tree_store(gtk.TreeStore):
	#constructor takes a xml_file and a set of arguments
	#as xml tag name and a python or gtk type
	def __init__(self, xml_file, update_xml, *args):
		#set up variables
		self.column_count = 0
		self.formatted_columns = []
		self.xml_file = xml_file
		self.update_xml = update_xml
		#if we got a filename, load the xml file
		if type(self.xml_file) == str:
			try:
				self.tree = etree.parse(xml_file)
			except:
				print("Invalid file " + xml_file)
			self.xml = self.tree.getroot()
		#if we got an etree Element, just copy it over to self.xml
		elif type(self.xml_file) == etree._Element:
			self.xml = self.xml_file
		#get the tag names of the root tag and the main element tag
		self.xml_root_tag = self.xml.tag
		self.xml_element_tag = self.xml[0].tag
		#separate the list of arguments into gtk-relevant
		#ones and xml-relevant ones
		self.xml_args = []
		self.gtk_args = []
		for children in args:
			if type(children) == type:
				self.gtk_args.append(children)
				self.column_count += 1
			else:
				self.xml_args.append(children)
		for i in range(self.column_count):
			self.formatted_columns.append(False)
		self.formatted_columns.append(False)
		self.gtk_args.append(object)
		#call gtk.TreeStore's init function
		gtk.TreeStore.__init__(self, *self.gtk_args)
		self.populate_treestore(None, self.xml)

	#---------------------------------------
	# New functions
	#---------------------------------------
	def format_column(self, column, formatting):
		#formatting should be the markup, minus brackets, for some pango
		#markup.  So if you want to make column 1 bold, call:
		#self.format_column(1, "b")
		if self.gtk_args[column] == str:
			self.foreach(self.format_func, (column, formatting))
			self.formatted_columns[column] = True

	def format_func(self, model, path, my_iter, user_data):
		#internally used to implement self.format_column
		(column, f) = user_data
		self.set_value(my_iter, column, "<" + f + ">" + self.get_value(my_iter, column) + "</" + f + ">")

	def populate_treestore(self, gtk_parent, xml_parent):
		#populate the treestore (for now, the only difference
		#it acknowledges is between int, bool, and str (which is what
		#it makes everything that isn't an int or bool)
		count = 0
		for children in self.xml_args:
			if len(xml_parent) != 0:
				for xml_children in xml_parent.findall(children):
					row = []
					count = 0
					if xml_children.tag == self.xml_element_tag:
						if self.gtk_args[count] != bool and self.gtk_args[count] != int:
							row.append(xml_children.text)
						elif gtk_args[count] == bool:
							if xml_children.text.lower() == "false":
								row.append(False)
							else:
								row.append(True)
						elif gtk_args[count] == int:
							row.append(int(xml_children.find(children).text))
					row.append(xml_children)
					parent = self.append(gtk_parent, row)
					self.populate_treestore(parent, xml_children)
				count += 1

	#---------------------------------------
	# Functions adapted from gtk.Treestore
	#---------------------------------------
	def append(self, parent, row=None):
		if self.update_xml == True:
			if parent != None:
				parent_element = self.get_value(parent, self.column_count)
			else:
				parent_element = self.xml
			element = SubElement(parent_element, self.xml_element_tag)
			for children in self.xml_args:
				SubElement(element, children)
		return gtk.TreeStore.append(self, parent, row)			

	def clear(self):
		#in addition to clearing the TreeStore, this function will
		#clear the xml tree unless update_xml is set to false
		for children in self.xml:
			self.xml.remove(children)
		gtk.TreeStore.clear()

	def get_value(self, my_iter, column):
		retval = gtk.TreeStore.get_value(self, my_iter, column)
		if self.formatted_columns[column] == True:
			retval = retval[3:]
			retval = retval[0:len(retval)-4]
		return retval

	def insert(self, parent, position, row=None):
		#in addition to inserting a row, this function will append
		#a new xml element to the tree and give it all the default
		#child elements unless update_xml is set to false
		if self.update_xml == True:
			parent_element = self.get_value(parent, self.column_count)
			element = parent_element.append(self.xml_element_tag)
			for children in self.xml_args:
				element.append(children)
		return gtk.TreeStore.insert(self, parent, position, row)

	def insert_after(self, parent, sibling, row=None):
		if self.update_xml == True:
			parent_element = self.get_value(parent, self.column_count)
			element = parent_element.append(self.xml_element_tag)
			for children in self.xml_args:
				element.append(children)
		return gtk.TreeStore.insert_after(self, sibling, row)

	def insert_before(self, parent, sibling, row=None):
		if self.update_xml == True:
			parent_element = self.get_value(parent, self.column_count)
			element = parent_element.append(self.xml_element_tag)
			for children in self.xml_args:
				element.append(children)
		return gtk.TreeStore.insert_before(self, sibling, row)

	def prepend(self, parent, row=None):
		if self.update_xml == True:
			parent_element = self.get_value(parent, self.column_count)
			element = parent_element.append(self.xml_element_tag)
			for children in self.xml_args:
				element.append(children)
		return gtk.TreeStore.prepend(self, parent, row)			

	def remove(self, my_iter):
		#removes both xml tree row and corresponding xml element
		#unless update_xml is set to False
		if self.update_xml == True:
			element = self.get_value(my_iter, self.column_count)
			element.getparent().remove(element)
		return gtk.TreeStore.remove(self, my_iter)

	def set(self, my_iter, *args):
		#sets the value of cells in the row my_iter points to and
		#updates the xml unless update_xml is set to False
		if self.update_xml == True:
			element = self.get_value(my_iter, self.column_count)
			for children in args:
				(column, value) = children
				element.text = value
		gtk.TreeStore.set(self, my_iter, args)

	def set_value(self, my_iter, column, value):
		#sets the value of the cell at my_iter, column to value and
		#updates the xml unless update_xml is set to False
		if self.update_xml == True:
			element = self.get_value(my_iter, self.column_count)
			element.text = value
		gtk.TreeStore.set_value(self, my_iter, column, value)

#This class automatically sets up a gtk.TreeView, including columns and
#CellRenderers, based on the gtk.TreeStore passed to it.  It will assume
#that a bool should be represented by a gtk.CellRendererToggle and that
#everything else should be represented by a gtk.CellRendererText, which
#will automatically have markup made available.  Columns will be stored in
#the list self.columns[], and CellRenderers will be stored in the list
#self.cell_renderers[], so that, to access the first column and the
#CellRenderer associated with it, you will only need to look to
#self.column[1] and self.cell_renderer[1], which should make for easy
#loops.  Any markup for text cells should be stored in the gtk.TreeStore.
#If any CellRenderTexts are made editable, they will already have been
#connected to the default callback: text_edited_cb, which should suffice.
class auto_tree_view(gtk.TreeView):
	def __init__(self, model):
		#set up variables
		self.cell_renderers = []
		self.columns = []
		self.my_model = model
		#call gtk.TreeView's __init__ function
		gtk.TreeView.__init__(self)
		self.set_model(model)
		#set up the columns and cell renderers:
		for i in range(self.my_model.get_n_columns()):
			self.columns.append(gtk.TreeViewColumn())
			column_type = self.my_model.get_column_type(i)
			if column_type == bool:
				self.cell_renderers.append(gtk.CellRendererToggle())
				self.cell_renderers[i].set_property("activatable", True)
				self.cell_renderers[i].connect("toggled", self.toggle_cb, i)
				self.columns[i].pack_start(self.cell_renderers[i])
				self.columns[i].add_attribute(self.cell_renderers[i], "active", i)
				self.append_column(self.columns[i])
			else:
				self.cell_renderers.append(gtk.CellRendererText())
				self.cell_renderers[i].connect("edited", self.text_edited_cb, i)
				self.columns[i].pack_start(self.cell_renderers[i])
				self.columns[i].add_attribute(self.cell_renderers[i], "markup", i)
				self.append_column(self.columns[i])
		if type(self.my_model) == xml_tree_store:
			self.remove_column(self.columns[len(self.columns) - 1])

	#---------------------------------------
	# Callbacks
	#---------------------------------------
	def text_edited_cb(self, text, path, new_text, column):
		#only do the stuff below if the text is valid.  By default, all text is valid.
		if self.text_edited_cb_validator(new_text) == True:
			#get a gtk.TreeIter from the path
			tree_iter = self.my_model.get_iter(path)
			self.my_model.set_value(path, column, new_text)

	def text_edited_cb_validator(self, new_text):
		#add code here if you want to validate new_text.  Return True if valid and
		#False if invalid.
		return True

	def toggle_cb(self, toggle, path, column):
		#get a gtk.TreeIter from the path
		tree_iter = self.my_model.get_iter(path)
		#I'm working in reverse here (in case the tree is sorted based on the bool column)
		#so False == True and True == False.
		if self.my_model.get_value(tree_iter, column) == False:
			#If we're checking the box, make sure parent boxes are checked, too.
			if self.my_model.iter_depth(tree_iter) != 0:
				parent = tree_iter
				for i in range(self.my_model.iter_depth(tree_iter)):
					parent = self.my_model.iter_parent(parent)
					self.my_model.set_value(parent, column, True)
		else:
			#If we're clearing the box, make sure child boxes are cleared, too.
			if self.my_model.iter_n_children(tree_iter) != 0:
				self.toggle_cb_recurse(tree_iter, column)
		#toggle the value of the TreeStore - we do this at the end just in case the tree is
		#sorted on this column.
		self.my_model.set_value(tree_iter, column, not(self.my_model.get_value(tree_iter, column)))
		#make sure the treestore displays the new value
		toggle.set_active(self.my_model.get_value(tree_iter, column))

	def toggle_cb_recurse(self, parent, column):
		for i in range(self.my_model.iter_n_children(parent)):
			child = self.my_model.iter_nth_child(i)
			toggle_cb_recurse(self, child, column)
			self.my_model.set_value(child, column, False)