прикреплённый файл:2013-03-01.drawtree.py к странице LecturesVMSH/Python/2013-03-01

   1 #!/usr/bin/env python
   2 # coding: utf
   3 '''
   4 Пользуясь стрелкой, алгоритмом рисования дерева и подмодулем pygame.font (или прилагаемой функцией Label()), написать программу визуализации дерева на PyGame.
   5 '''
   6 
   7 import pygame, random
   8 from math import *
   9 from label import Label
  10 
  11 pygame.init()
  12 Size=(1024,768)
  13 Scr=pygame.display.set_mode(Size)
  14 R=12
  15 Arrows={}
  16 
  17 def length(m,M):
  18     '''Расстояние от m до M'''
  19     return sqrt((m[0]-M[0])**2+(m[1]-M[1])**2)
  20 
  21 def getarrow(size=9, color=(0,0,0), back=(255,255,255)):
  22     '''Ищет стрелку размера size и цвета color на складе,
  23     если такой нет — добавляет её туда с прозрачным цветом фона back'''
  24     if (color[:3],size) not in Arrows:
  25         Arrows[color[:3],size]=pygame.Surface((size,size))
  26         Arrows[color[:3],size].fill(back)
  27         pygame.draw.polygon(Arrows[color[:3],size], color, ((size-1,0),(0,size/2),(size-1,size-1),(size/2,size/2)))
  28         Arrows[color[:3],size].set_colorkey(back)
  29     return Arrows[color[:3],size]
  30 
  31 def arrow(scr,m,M,R,color=(0,0,0),size=9):
  32     '''Рисует стрелку размера size цвета color из m в M с отсупом в R от M и m'''
  33     ar=getarrow(size,color)
  34     L=length(m,M)
  35     l=L-2*R-size/2
  36     n=m[0]+(M[0]-m[0])*R/L,m[1]+(M[1]-m[1])*R/L
  37     N=m[0]+(M[0]-m[0])*(R+l)/L,m[1]+(M[1]-m[1])*(R+l)/L
  38     arr=pygame.transform.rotate(ar,degrees(atan2(M[1]-n[1],n[0]-M[0])))
  39     pygame.draw.line(scr,color,n,N)
  40     scr.blit(arr,arr.get_rect(center=N))
  41 
  42 def drawtree(Tr,root,pos,dpos=R*2):
  43     '''Рисовать дерево Tr:
  44     Рисовать узел root, рисовать поддеревья со смещением dpos, рисовать стрелки'''
  45     Label(Scr,str(root),pos,R,(0,0,0),centered=True)
  46     # По ходу будем вычислять прямоугольник, описаный вокруг дерева
  47     rect=pygame.Rect(pos,(0,0))
  48     for subr in Tr[root]:
  49         newrect=drawtree(Tr,subr,(rect.right,pos[1]+dpos),dpos)
  50         arrow(Scr,pos,newrect.topleft,R)
  51         rect.union_ip(newrect)
  52     if Tr[root]:
  53         return rect
  54     else:
  55         return pygame.Rect(pos,(dpos,dpos))
  56 
  57 MAXN=30
  58 Tree=[(random.randrange(i/4,i),i) for i in xrange(1,MAXN)]
  59 DTree={}
  60 for i in xrange(MAXN): DTree[i]=[]
  61 for a,b in Tree: DTree[a].append(b)
  62 print DTree
  63 
  64 CanExit, NeedRedraw = False, True
  65 while not CanExit:
  66     ev=pygame.event.wait()
  67     if ev.type == pygame.QUIT:
  68         CanExit=True
  69     if NeedRedraw:
  70         Scr.fill(pygame.Color("wheat"))
  71         drawtree(DTree,0,(R,R),R*4)
  72         NeedRedraw=False
  73     pygame.display.flip()