indentation and imports cleaned up
[p2pool.git] / p2pool / draw.py
1 import pygame
2 import time
3 import hashlib
4 import math
5 import StringIO
6 from PIL import Image
7
8 from p2pool.util.vector import v
9
10 @apply
11 class color(object):
12     def __getattr__(self, name):
13         res = pygame.Color(name)
14         setattr(self, name, res)
15         return res
16
17 def get_uniform(bound, *data):
18     x = int(hashlib.sha256(repr(data)).hexdigest(), 16)
19     return x % bound
20
21 def get_pos(share, t, d):
22     x = 5 + get_uniform(d.get_width() - 10, share.hash, "pos")
23     y = d.get_height() - (t - share.time_seen)*10
24     return v(x, y)
25
26 def get_color(data):
27     return [get_uniform(256, data, x) for x in "rgb"]
28
29 def perp_and_normalize_to((dx, dy), d):
30     m = math.sqrt(dx**2 + dy**2)
31     return v(-dy/m*d, dx/m*d)
32
33 def go(share, tracker, t, d):
34     #c = color.green if share.peer is None else color.red
35     c = get_color(share.new_script)
36     pos = get_pos(share, t, d)
37     pygame.draw.circle(d, c, pos.rounded, 5)
38     if share.previous_hash in tracker.shares:
39         previous_share = tracker.shares[share.previous_hash]
40         previous_pos = get_pos(previous_share, t, d)
41         vec_to_previous = previous_pos - pos
42         pygame.draw.polygon(d, c, [
43             (pos + perp_and_normalize_to(vec_to_previous, 5)).rounded,
44             (pos + perp_and_normalize_to(vec_to_previous, -5)).rounded,
45             previous_pos.rounded,
46         ])
47     if share.peer is None:
48         pygame.draw.circle(d, c, pos.rounded, 10, 2)
49     for child_hash in tracker.reverse_shares.get(share.hash, set()):
50         go(tracker.shares[child_hash], tracker, t, d)
51
52 def get(tracker, best):
53     d = pygame.Surface((400, 600), 32)
54     if tracker.get_height(best) >= 100:
55         t = time.time()
56         start = tracker.get_nth_parent_hash(best, 100)
57         d.fill((0, 0, 0))
58         go(tracker.shares[start], tracker, t, d)
59     f = StringIO.StringIO()
60     Image.fromstring("RGB", d.get_size(), pygame.image.tostring(d, "RGB")).save(f, "png")
61     return f.getvalue()