assembly_handbook/ahb_techdraw_extensions.py

import FreeCAD as App
import FreeCADGui as Gui

from PySide.QtCore import QTimer

import TechDraw, TechDrawGui
from PySide import QtGui, QtCore
TDG = TechDrawGui

class CursorItem(QtGui.QGraphicsItem):
	def __init__(self, parent = None, view = None):
		super().__init__(parent)
		self.Type = QtGui.QGraphicsItem.UserType + 501
		self.setZValue(500)
		self.margin = 10.0
		self.size = 100.0
		self.view = view
		
	def onViewPosChange(self, callback):
		self.viewPosChangeCallback = callback

	def boundingRect(self):
		return QtCore.QRectF(-self.size/2 - self.margin, -self.size/2 - self.margin,
					  self.size + 2.0*self.margin, self.size + 2.0*self.margin)

	def paint(self, painter, option, widget):
		#print("paint")
		painter.setBrush(QtCore.Qt.darkRed)
		#painter.drawRoundedRect(-100, -100, 200, 200, 50, 50)
		h = self.size/2.0
		painter.drawLine(-h, 0, h, 0)
		painter.drawLine(0, -h, 0, h)
		
	def mousePressEvent(self, event):
		#print('mouse press', event)
		self.startMovePos = event.pos()
		
	def mouseMoveEvent(self, event):
		#print('mouse move', event)
		offset = event.pos() - self.startMovePos
		self.moveBy(offset.x(), offset.y())
		
	def mouseReleaseEvent(self, event):
		#print('mouse release', event)
		#print('new pos', self.x(), self.y())
		if self.viewPosChangeCallback is not None:
			self.viewPosChangeCallback(self.getViewPos())
	
	def getViewPos(self):
		scale = self.view.Scale * 10.0
		return App.Vector(self.x() / scale, -self.y() / scale)
	
	def setViewPos(self, p):
		scale = self.view.Scale * 10.0
		self.setPos(p.x * scale, -p.y * scale)

class TechDrawExtensions:
	views_to_repaint = {}
	
	view_cursors = {}
	
	updating_balloon = False
	
	edited_view = None
	
	def __init__(self):
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		workbench.docObserver.onObjectTypeChanged('balloon_changed', 'TechDraw::DrawViewBalloon', lambda obj, prop: self.onBalloonChanged(obj, prop))
		workbench.docObserver.onObjectTypeSelected('balloon_selected', 'TechDraw::DrawViewBalloon', lambda operation, obj, sub, point: self.onBalloonSelected(operation, obj, sub, point))
	
	def repaint(self, view):
		self.views_to_repaint[view] = True
		QTimer.singleShot(100, self._do_repaint)
		
	def _do_repaint(self):
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		to_repaint = self.views_to_repaint.keys()
		self.views_to_repaint = {}
		
		for view in to_repaint:
			#print("Repainting " + view.Name)
			
			selected_balloons = []
			for obj in Gui.Selection.getSelection():
				if obj.TypeId == 'TechDraw::DrawViewBalloon' and obj.SourceView == view and 'Assembly_handbook_PartName' in obj.PropertiesList:
					selected_balloons.append(obj)
					
			is_first_part = True
			
			if len(selected_balloons) == 0:
				selected_balloons.append(None)
			
			for balloon in selected_balloons:
				if balloon is not None:
					doc = balloon.Document
					partLink = doc.getObject(balloon.Assembly_handbook_PartName)
					
					part_view = workbench.partsCache.getPart2DView(view, partLink)
					
					center = self.computePartCenter(view, partLink)
				
				# iterate edges of actual view and highlight matching edges
				for edgeIdx in range(10000):
					hasEdge = False
					try:
						edge = view.getEdgeByIndex(edgeIdx)
						hasEdge = True
					except:
						pass
					if not hasEdge:
						break
					
					# reset edge format
					if is_first_part:
						view.formatGeometricEdge(edgeIdx,1,0.25,0,True)
					
					if balloon is not None and (not hasattr(edge.Curve, 'Degree') or edge.Curve.Degree == 1) and len(edge.Vertexes) == 2:
						edgeData = [
							edge.Vertexes[0].X - center.x,
							edge.Vertexes[0].Y - center.y,
							edge.Vertexes[1].X - center.x,
							edge.Vertexes[1].Y - center.y
						]
						v0 = App.Vector(edgeData[0], edgeData[1])
						v1 = App.Vector(edgeData[2], edgeData[3])
						
						for line in part_view.cached_lines:
							l0 = App.Vector(line[0], line[1])
							l1 = App.Vector(line[2], line[3])
							#d = abs(edgeData[0] - line[0]) + abs(edgeData[1] - line[1]) + abs(edgeData[2] - line[2]) + abs(edgeData[3] - line[3])
							d = v0.distanceToLineSegment(l0, l1).Length + v1.distanceToLineSegment(l0, l1).Length
							if d < 0.01:
								view.formatGeometricEdge(edgeIdx,1,0.25,(0,0.85,0),True)
				
				is_first_part = False
			
			view.requestPaint()
			
			cursor = self.view_cursors.get(view, None)
			if cursor is None:
				cursor = CursorItem(view = view)
				TDG.addQGIToView(view, cursor);
				cursor.onViewPosChange(lambda new_pos: self.onCursorMoved(view, new_pos))
				self.view_cursors[view] = cursor
				
			if selected_balloons[0] is None:
				cursor.setVisible(False)
			else:
				cursor.setViewPos(App.Vector(selected_balloons[0].OriginX, selected_balloons[0].OriginY))
				cursor.setVisible(True)
				
	def refreshView(self, view):
		doc = view.Document
		page = self.getViewPage(view)
		for balloon in page.Views:
			if balloon.TypeId == 'TechDraw::DrawViewBalloon' and "Assembly_handbook_PartName" in balloon.PropertiesList and balloon.SourceView == view:
				obj = doc.getObject(balloon.Assembly_handbook_PartName)
				balloonColor = (0.0, 0.0, 0.0)
				if obj is None or not obj in view.XSource:
					balloonColor = (1.0, 0.0, 0.0)
				balloon.ViewObject.Color = balloonColor
		
				
	def toggleEditViewSourceParts(self, view):
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		button = self.getToolbarButton('AHB_view_edit_source_parts')
		
		if self.edited_view is None:
			workbench.docLinkObserver.select_link_mode = True
			self.edited_view = view
			
			button.setChecked(True)
			button.setText('End source parts edition')
			
			Gui.Selection.clearSelection()
			for obj in view.XSource:
				Gui.Selection.addSelection(obj)
		else:
			workbench.docLinkObserver.select_link_mode = False
			
			Gui.Selection.clearSelection()
			Gui.Selection.addSelection(self.edited_view)
			self.edited_view = None
			
			button.setChecked(False)
			button.setText('Edit view source parts')
			
	def editViewSourceParts(self, parts, add):
		if self.edited_view is None: return
		xsource = self.edited_view.XSource
		modified = False
		for part in parts:
			if (part in xsource) != add:
				if add:
					xsource.append(part)
				else:
					xsource.remove(part)
				modified = True
		if modified:
			self.edited_view.XSource = xsource
		
	def getToolbarButton(self, buttonName):
		mainwin = Gui.getMainWindow()
		toolbar = None
		for tb in mainwin.findChildren(QtGui.QToolBar):
			if tb.objectName()=='Assembly Handbook':
				toolbar = tb
				
		button = None
		if toolbar is not None:
			for action in toolbar.actions():
				if action.objectName() == buttonName:
					button = action
		
		return button
				
	def onCursorMoved(self, view, new_pos):
		if len(Gui.Selection.getSelection()) == 0: return
		balloon = Gui.Selection.getSelection()[0]
		if balloon.TypeId != 'TechDraw::DrawViewBalloon' or not 'Assembly_handbook_PartName' in balloon.PropertiesList: return
		if balloon.SourceView != view: return
		
		balloon.OriginX = new_pos.x
		balloon.OriginY = new_pos.y
		
		obj = balloon.Document.getObject(balloon.Assembly_handbook_PartName)
		view = balloon.SourceView
		center = self.computePartCenter(view, obj)
		
		balloon.Assembly_handbook_OriginOffsetX = new_pos.x - center.x
		balloon.Assembly_handbook_OriginOffsetY = new_pos.y - center.y
	
	def onBalloonSelected(self, operation, balloon, sub, point):
		import time
		
		#print(operation, obj.Name, sub, point)
		if "Assembly_handbook_PartName" in balloon.PropertiesList:
			#print(operation + " " + balloon.Name)
			view = balloon.SourceView
			self.repaint(view)
			
	def onBalloonChanged(self, obj, prop):
		# Avoid reentry
		if self.updating_balloon:
			return
		
		#print('Balloon changed: ' + obj.Name + '.' + prop)
		if prop == 'Y' and "Assembly_handbook_PartName" in obj.PropertiesList:
			self.updating_balloon = True
			self.updateBalloon(obj)
			self.updating_balloon = False
	
	def updateBalloon(self, balloon):
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		view = balloon.SourceView
		doc = view.Document
		obj = doc.getObject(balloon.Assembly_handbook_PartName)
		
		objectCenterView = workbench.techDrawExtensions.computePartCenter(view, obj)
		
		balloon.OriginX = objectCenterView.x + balloon.Assembly_handbook_OriginOffsetX
		balloon.OriginY = objectCenterView.y + balloon.Assembly_handbook_OriginOffsetY
		
		balloon.Text = obj.LinkedObject.Document.Name if obj.TypeId == 'App::Link' else obj.Name
		balloon.ViewObject.Font = 'DejaVu Sans'
		balloon.ViewObject.Fontsize = 4
		balloon.BubbleShape = 'Inspection'
		balloon.EndTypeScale = 4
		
	def getActivePage(self):
		activeView = Gui.activeView()
		if activeView is None: return None
		activePage = activeView.getPage() if hasattr(activeView, 'getPage') else None
		return activePage
				
	def getViewPage(self, view):
		for obj in view.InList:
			if obj.TypeId == 'TechDraw::DrawPage':
				if view in obj.Views:
					return  obj
		return None
	
	def forceRedrawPage(self, page, callback = None):
		for view in page.Views:
			if view.TypeId == 'TechDraw::DrawViewPart' and 'Assembly_handbook_PreviousStepView' in view.PropertiesList:
				self.refreshView(view) 
		
		if page.KeepUpdated:
			for view in page.Views:
				if view.TypeId != 'TechDraw::DrawViewPart':
					view.recompute()
			for view in page.Views:
				if view.TypeId == 'TechDraw::DrawViewPart':
					view.recompute() 
			if callback is not None:
				callback()
		else:
			page.KeepUpdated = True
			def restoreKeepUpdated():
				page.KeepUpdated = False
				if callback is not None:
					callback()
			QTimer.singleShot(10, restoreKeepUpdated)
	
	def computePartCenter(self, view, obj):
		if obj.TypeId == 'App::Link':
			partLink = obj
			objectCenterWorld = partLink.LinkPlacement.Matrix.multiply(partLink.LinkedObject.Shape.CenterOfGravity)
		else:
			objectCenterWorld = obj.Shape.CenterOfGravity
		
		vertId = view.makeCosmeticVertex3d(objectCenterWorld)
		vert = view.getCosmeticVertex(vertId)
		objectCenterView = vert.Point
		view.removeCosmeticVertex(vertId)
		return objectCenterView