assembly_handbook/ahb_cmd_view_annotate.py

import math
import FreeCADGui as Gui
import FreeCAD as App

import ahb_utils

class AHB_View_Annotate:
	def GetResources(self):
		return {"MenuText": "Annotate view",
				"ToolTip": "Annotates a TechDraw view with object names",
				"Pixmap": ""
				}

	def IsActive(self):
		return ahb_utils.getCurrentView() is not None

	def Activated(self):
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		if len(Gui.Selection.getSelection()) == 0:
			page = workbench.techDrawExtensions.getActivePage()
			workbench.techDrawExtensions.refreshOverlays(page)
			return
		
		if len(Gui.Selection.getSelection()) != 1:
			raise Exception("Please select exactly one TechDraw view")

		view = Gui.Selection.getSelection()[0]
		if view.TypeId == 'TechDraw::DrawPage':
			workbench.techDrawExtensions.refreshOverlays(view)
			return
		elif view.TypeId != 'TechDraw::DrawViewPart':
			raise Exception("Selected object is not a TechDraw view")
		
		overlay_view = workbench.techDrawExtensions.getOverlayView(view)
		
		doc = view.Document
		
		page = workbench.techDrawExtensions.getViewPage(view)
		if page is None:
			raise Exception("Can't find page in which the selected view is located")
		
		def list_sub_parts(parts):
			result = []
			for part in parts:
				if part.TypeId == 'App::Link':
					result.append(part)
				elif part.TypeId == 'Part::FeaturePython' and hasattr(part, 'LinkedObject'): # variant link
					result.append(part)
				
				if hasattr(part, 'Group'):
					result.extend(list_sub_parts(part.Group))
			
			return result	
		
		all_parts = list_sub_parts(view.XSource)
		
		# Remove balloons referencing missing objects
		for balloon in page.Views:
			if balloon.TypeId == 'TechDraw::DrawViewBalloon' and "Assembly_handbook_Source" in balloon.PropertiesList:
				if balloon.SourceView != view and balloon.SourceView != overlay_view: continue
				partLink = workbench.techDrawExtensions.getBalloonSourcePart(balloon)
				if partLink is None or partLink not in all_parts:
					ref_name = "<no ref>"
					try:
						ref_name = balloon.Assembly_handbook_Source[1]
					except:
						pass
					print(balloon.Name + " references missing object " + ref_name + ", removing balloon")
					doc.removeObject(balloon.Name)
		
		for partLink in view.XSource:
			workbench.techDrawExtensions.add_or_update_balloon(view, partLink, '')
			
		workbench.techDrawExtensions.refreshOverlays(page)

		regroupedBalloons = self.RegroupNearestSimilarBalloons(balloonsCreated)
		self.PlaceBalloonsInCircle(regroupedBalloons)

	def CalculatePointsCenter(self, balloons):
		totalX = 0
		totalY = 0
		for balloon in balloons:
			realBalloon = balloon[0] if type(balloon) is list else balloon
			totalX = totalX + int(realBalloon.OriginX)
			totalY = totalY + int(realBalloon.OriginY)
		return App.Base.Vector2d(totalX / len(balloons), totalY / len(balloons))

	def IsSimilarBalloonNear(self, balloonA, balloonB):
		MAX_DISTANCE_BETWEEN_REGROUPED_BALLOONS = 50
		if balloonA.Text == balloonB.Text:
			pos = App.Base.Vector2d(balloonA.OriginX, balloonA.OriginY)
			dist = pos.distance(App.Base.Vector2d(balloonB.OriginX, balloonB.OriginY))
			return dist < MAX_DISTANCE_BETWEEN_REGROUPED_BALLOONS
		else:
			return False

	def RegroupNearestSimilarBalloons(self, balloons):
		regroupedBalloons = []
		for balloon in balloons:
			nearestBalloons = []
			for otherBalloon in balloons:
				if otherBalloon != balloon and self.IsSimilarBalloonNear(balloon, otherBalloon):
					nearestBalloons.append(otherBalloon)
					balloons.remove(otherBalloon)
			if len(nearestBalloons) == 0:
				regroupedBalloons.append(balloon)
			else:
				nearestBalloons.append(balloon)
				regroupedBalloons.append(nearestBalloons)
		return regroupedBalloons

	def PlaceBalloonsInCircle(self, balloons):
		center = self.CalculatePointsCenter(balloons)
		nbBalloons = len(balloons)
		balloonPosStep = (math.pi * 2) / nbBalloons
		for i in range(nbBalloons):
			balloonPos = App.Base.Vector2d(center.x + 600 * math.cos(balloonPosStep * i), center.y + 600 * math.sin(balloonPosStep * i))
			# Find nearest arrow to avoid arrow crossing each other
			smallestDistance = 0
			balloonToUse = None
			for balloon in balloons:
				realBalloon = balloon[0] if type(balloon) is list else balloon
				dist = balloonPos.distance(App.Base.Vector2d(realBalloon.OriginX, realBalloon.OriginY))
				if smallestDistance == 0 or dist < smallestDistance:
					smallestDistance = dist
					balloonToUse = balloon

			if balloonToUse is not None:
				balloons.remove(balloonToUse)
				if type(balloonToUse) is list:
					for realBalloon in balloonToUse:
						realBalloon.X = balloonPos.x
						realBalloon.Y = balloonPos.y
				else:
					balloonToUse.X = balloonPos.x
					balloonToUse.Y = balloonPos.y

from ahb_command import AHB_CommandWrapper
AHB_CommandWrapper.addGuiCommand('AHB_view_annotate', AHB_View_Annotate())