scripts/parse/ctree_draw.py

#! /usr/bin/python
# -*- coding: utf-8 -*-

"""	
	This module contains classes to draw constituency tree.
	
	Version 0.1											(12-Feb-2015)
	ConstTreeDrawer, CTDCairoContext, DrawConstTree and DrawConstNode are added.
	
"""

import cairo
import parse


class ConstTreeDrawer():
	'''
	Class for drawing constituency tree
	'''
	
	def __init__(self, pOutFilename, pFontSize = 10):
		'''
		Constructor
		'''
		
		self.tree = None
		
		# drawing attributes
		
		self.width = 0
		self.height = 0
		
		self.branchHeight = 50
		# number of spaces between terminals
		self.terminalSpaceCount = 4
		self.margins = 10
		
		self.surface = cairo.PDFSurface(pOutFilename, self.width, self.height)
		self.context = CTDCairoContext(self.surface)
		
		self.context.setFontSize(pFontSize)
		self.context.setRGB(0, 0, 0)
		self.context.setLineWidth(.25)
		
	
	
	def _createRoot(self):
		'''
		Creates and returns the root node
		'''
		
		return ConstNodeDraw()
		
	
	
	def draw(self, pCTree):
		'''
		Draws a given constituency tree 
		
		The tree can be provided as a string in bracketing format or as
		parse.ConstTree type. The string bracketing format can contain 
		color information or can be plain. The ConstTree object can be of
		either ConstTree or DrawConstTree type.
		'''
		
		# preparing the tree
		
		if isinstance(pCTree, parse.ConstTree):
			self.tree = DrawConstTree()
			self.tree.loadPTBTree(pCTree.getPTBFormat(), pflgExpandTerminal = True)
		elif isinstance(pCTree, basestring):
			self.tree = DrawConstTree()
			self.tree.loadPTBTree(pCTree, pflgExpandTerminal = True)
		else:
			raise Exception("Unknown constituency tree format!")
		
		
		# presetting of drawing attributes of the tree
		self.tree.presetDrawAttrs(self.context)
		
		# moving the position to the top left corner of the draw area
		self.context.move_to(self.margins, self.height + self.margins)
		
		# setting the surface size
		self.setSurfaceSize()
		
		# actual drawing
		#self.tree.draw(pCanvas = self.context, pBranchHeight = self.branchHeight, pTerminalSpaceCount = self.terminalSpaceCount)
		self.context.show_text('  '.join(self.tree.getTerminals()))
		
		# finishing the draw by applying it on the surface
		self.surface.finish()
		
	
	
	def setSurfaceSize(self):
		'''
		Sets the size of the drawing surface
		'''
		
		# caculating the tree area
		self.width = self.tree.getDrawWidth() + (2 * self.margins)
		self.height = self.tree.getDrawHeight + (2 * self.margins)
		self.surface.set_size(self.width, self.height)
		
		
	
	


class CTDCairoContext(cairo.Context):
	'''
	A customized Cairo Context class for drawing constituency tree
	'''
	
	
	def __init__(self, pCairoSurface):
		'''
		Constructor
		'''
		
		cairo.Context.__init__(self, pCairoSurface)
		
		self.fontSize = None
		
	
	
	def setFontSize(self, pFontSize):
		'''
		Sets the font size of the context
		'''
		
		self.set_font_size(pFontSize)
		
		self.fontSize = pFontSize
		
	
	
	def getFontSize(self):
		'''
		Returns the value of the font size
		'''
		
		return self.fontSize
		
	
	
	def setRGB(self, pR, pG, pB):
		'''
		Sets the RGB color of the context
		'''
		
		self.set_source_rgb(pR, pG, pB)
		
		self.RGB = (pR, pG, pB)
		
	
	
	def getRGB(self):
		'''
		Returns the value of the RGB color of the context
		'''
		
		return self.RGB
		
	
	
	def setLineWidth(self, pLineWidth):
		'''
		Sets the line width of the context
		'''
		
		self.set_line_width(pLineWidth)
		
		self.lineWidth = pLineWidth
		
	
	
	def getLineWidth(self):
		'''
		Returns the value of the line width of the context
		'''
		
		return self.lineWidth
		
	


class DrawConstTree(parse.ConstTree):
	'''
	Class for constituency tree with drawing features
	'''
	
	def __init__(self):
		'''
		Constructor
		'''
		
		parse.ConstTree.__init__(self)
		
		self.drawWidth = 0
		self.drawHeight = 0
		
	
	
	def _createRoot(self):
		'''
		Creates and returns the root node
		'''
		
		return DrawConstNode()
		
	
	
	def presetDrawAttrs(self, pCTDContext):
		'''
		Presets the drawing attributes of the tree and its nodes 
		'''
		
		self.root.presetDrawAttrs(pCTDContext)
		
		self.drawWidth 
		self.drawHeight
		
	
	
	def draw(self, pCanvas, pBranchHeight, pTerminalSpaceCount):
		'''
		Draws the tree on the given canvas with the provided drawing 
		attributes
		'''
		
		self.root.draw(pCanvas = pCanvas, pBranchHeight = pBranchHeight, pTerminalSpaceCount = pTerminalSpaceCount, pLevel = 0)
		
	
	


class DrawConstNode(parse.ConstNode):
	'''
	Class for constituency tree node with drawing features
	'''
	
	def __init__(self):
		'''
		Constructor
		'''
		
		parse.ConstNode.__init__(self)
		
		self.textWidth = 0
		self.textHeight = 0
		self.drawWidth = 0
		self.drawHeight = 0
		
	
	
	def _getNewNode(self):
		'''
		Creates and returns a node
		'''
		
		return DrawConstNode()
		
	
	
	def presetDrawAttrs(self, pCTDContext):
		'''
		Presets the drawing attributes of the nodes and the subtree itself
		'''
		
		if self.isTerminal():
			######
			vTextExtent = pCTDContext.text_extents(self.surface)
			self.drawWidth = vTokensExtent[2]
			self.drawHeight = vTokensExtent[3] 
			######
			
			return self.drawWidth, self.drawHeight
		else:
			vTextExtent = pCTDContext.text_extents(self.getLabel())
			self.textWidth = vTokensExtent[2]
			self.textHeight = vTokensExtent[3] 
			
			vChildrenDrawWidth = 0
			vChildrenDrawHeight = 0
			for vChild in self.getChildren():
				vChildDrawWidth, vChildDrawHeight = vChild.presetDrawAttrs(pCTDContext)
				vChildrenDrawWidth += vChildDrawWidth
				vChildrenDrawHeight += vChildDrawHeight
			
			if self.textWidth > 
		
	
	
	def draw(self, pCanvas, pBranchHeight, pTerminalSpaceCount, pLevel):
		'''
		Draws the subtree rooted at the node on the given canvas with the
		provided drawing attributes
		
		pLevel identifies the level or depth of the node in the tree. The 
		root node level is 0.
		'''
		
		# depth first drawing: children first
		
		for vChild in self.getChildren():
			vChild.draw(pCanvas = pCanvas, pBranchHeight = pBranchHeight, pTerminalSpaceCount = pTerminalSpaceCount, pLevel = pLevel + 1)
		
		# setting text position
		
		vText = (' ' * pTerminalSpaceCount).join(self.getTerminals())
		pCanvas.show_text(vText)
		print vText
		
		pCanvas.rel_move_to(0, pBranchHeight * pLevel)
		
		# writing the text
		
		if self.isTerminal():
			pCanvas.show_text(vText)
		else:
			pCanvas.show_text(self.getSynTag())