Permalink
Fetching contributors…
Cannot retrieve contributors at this time
337 lines (277 sloc) 12.1 KB
"""A basic playground. Most interesting function is draw a shape, basically
move the mouse as you want and pymunk will approximate a Poly shape from the
drawing.
"""
__docformat__ = "reStructuredText"
import pygame
from pygame.locals import *
from pygame.color import *
import pymunk as pm
from pymunk import Vec2d
import pymunk.util as u
#TODO: Clean up code
COLLTYPE_DEFAULT = 0
COLLTYPE_MOUSE = 1
class PhysicsDemo:
def flipyv(self, v):
return int(v.x), int(-v.y+self.h)
def __init__(self):
self.running = True
### Init pygame and create screen
pygame.init()
self.w, self.h = 600,600
self.screen = pygame.display.set_mode((self.w, self.h))
self.clock = pygame.time.Clock()
### Init pymunk and create space
self.space = pm.Space()
self.space.gravity = (0.0, -900.0)
### Walls
self.walls = []
self.create_wall_segments( [(100, 50), (500, 50)] )
## Balls
#balls = [createBall(space, (100,300))]
self.balls = []
### Polys
self.polys = []
h = 10
for y in range(1,h):
#for x in range(1, y):
x = 0
s = 10
p = Vec2d(300,40) + Vec2d(0, y*s*2)
self.polys.append(self.create_box(p, size=s, mass = 1))
self.run_physics = True
### Wall under construction
self.wall_points = []
### Poly under construction
self.poly_points = []
self.shape_to_remove = None
self.mouse_contact = None
def draw_helptext(self):
font = pygame.font.Font(None, 16)
text = ["LMB: Create ball"
,"LMB + Shift: Create box"
,"RMB on object: Remove object"
,"RMB(hold) + Shift: Create polygon, release to finish (we be converted to a convex hull of the points)"
,"RMB + Ctrl: Create wall, release to finish"
,"Space: Stop physics simulation"
,"k: Spawn a bunch of blocks"
,"f: Fire a ball from the top left corner"
]
y = 5
for line in text:
text = font.render(line, 1,THECOLORS["black"])
self.screen.blit(text, (5,y))
y += 10
def create_ball(self, point, mass=1.0, radius=15.0):
moment = pm.moment_for_circle(mass, 0.0, radius)
ball_body = pm.Body(mass, moment)
ball_body.position = Vec2d(point)
ball_shape = pm.Circle(ball_body, radius)
ball_shape.friction = 1.5
ball_shape.collision_type = COLLTYPE_DEFAULT
self.space.add(ball_body, ball_shape)
return ball_shape
def create_box(self, pos, size = 10, mass = 5.0):
box_points = [(-size, -size), (-size, size), (size,size), (size, -size)]
return self.create_poly(box_points, mass = mass, pos = pos)
def create_poly(self, points, mass = 5.0, pos = (0,0)):
moment = pm.moment_for_poly(mass, points)
#moment = 1000
body = pm.Body(mass, moment)
body.position = Vec2d(pos)
shape = pm.Poly(body, points)
shape.friction = 0.5
shape.collision_type = COLLTYPE_DEFAULT
self.space.add(body, shape)
return shape
def create_wall_segments(self, points):
"""Create a number of wall segments connecting the points"""
if len(points) < 2:
return []
points = list(map(Vec2d, points))
for i in range(len(points)-1):
v1 = Vec2d(points[i].x, points[i].y)
v2 = Vec2d(points[i+1].x, points[i+1].y)
wall_body = pm.Body(body_type=pm.Body.STATIC)
wall_shape = pm.Segment(wall_body, v1, v2, .0)
wall_shape.friction = 1.0
wall_shape.collision_type = COLLTYPE_DEFAULT
self.space.add(wall_shape)
self.walls.append(wall_shape)
def run(self):
while self.running:
self.loop()
def draw_ball(self, ball):
body = ball.body
v = body.position + ball.offset.cpvrotate(body.rotation_vector)
p = self.flipyv(v)
r = ball.radius
pygame.draw.circle(self.screen, THECOLORS["blue"], p, int(r), 2)
def draw_wall(self, wall):
body = wall.body
pv1 = self.flipyv(body.position + wall.a.cpvrotate(body.rotation_vector))
pv2 = self.flipyv(body.position + wall.b.cpvrotate(body.rotation_vector))
pygame.draw.lines(self.screen, THECOLORS["lightgray"], False, [pv1,pv2])
def draw_poly(self, poly):
body = poly.body
ps = [p.rotated(body.angle) + body.position for p in poly.get_vertices()]
ps.append(ps[0])
ps = list(map(self.flipyv, ps))
if u.is_clockwise(ps):
color = THECOLORS["green"]
else:
color = THECOLORS["red"]
pygame.draw.lines(self.screen, color, False, ps)
def draw(self):
### Clear the screen
self.screen.fill(THECOLORS["white"])
### Display some text
self.draw_helptext()
### Draw balls
for ball in self.balls:
self.draw_ball(ball)
### Draw walls
for wall in self.walls:
self.draw_wall(wall)
### Draw polys
for poly in self.polys:
self.draw_poly(poly)
### Draw Uncompleted walls
if len(self.wall_points) > 1:
ps = [self.flipyv(Vec2d(p)) for p in self.wall_points]
pygame.draw.lines(self.screen, THECOLORS["gray"], False, ps, 2)
### Uncompleted poly
if len(self.poly_points) > 1:
ps = [self.flipyv(Vec2d(p)) for p in self.poly_points]
pygame.draw.lines(self.screen, THECOLORS["red"], False, ps, 2)
### Mouse Contact
if self.mouse_contact is not None:
p = self.flipyv(self.mouse_contact)
pygame.draw.circle(self.screen, THECOLORS["red"], p, 3)
### All done, lets flip the display
pygame.display.flip()
def loop(self):
for event in pygame.event.get():
if event.type == QUIT:
self.running = False
elif event.type == KEYDOWN and event.key == K_ESCAPE:
self.running = False
elif event.type == KEYDOWN and event.key == K_p:
pygame.image.save(self.screen, "playground.png")
elif event.type == MOUSEBUTTONDOWN and event.button == 1: # LMB
if pygame.key.get_mods() & KMOD_SHIFT:
p = self.flipyv(Vec2d(event.pos))
self.polys.append(self.create_box(pos = p))
else:
#t = -10000
p = self.flipyv(Vec2d(event.pos))
self.balls.append(self.create_ball(p))
elif event.type == MOUSEBUTTONDOWN and event.button == 3: #RMB
if pygame.key.get_mods() & KMOD_SHIFT:
pass
elif pygame.key.get_mods() & KMOD_CTRL:
p = self.flipyv(Vec2d(event.pos))
self.wall_points.append(p)
elif self.shape_to_remove is not None:
self.balls = list(filter(lambda a: a != self.shape_to_remove, self.balls))
self.walls = list(filter(lambda a: a != self.shape_to_remove, self.walls))
self.polys = list(filter(lambda a: a != self.shape_to_remove, self.polys))
self.space.remove(self.shape_to_remove.body, self.shape_to_remove)
elif event.type == KEYUP and event.key in (K_RCTRL, K_LCTRL):
### Create Wall
self.create_wall_segments(self.wall_points)
self.wall_points = []
elif event.type == KEYUP and event.key in (K_RSHIFT, K_LSHIFT):
### Create Polygon
if len(self.poly_points) > 0:
self.poly_points = u.reduce_poly(self.poly_points, tolerance=5)
if len(self.poly_points) > 2:
self.poly_points = u.convex_hull(self.poly_points)
if not u.is_clockwise(self.poly_points):
self.poly_points.reverse()
center = u.calc_center(self.poly_points)
self.poly_points = u.poly_vectors_around_center(self.poly_points)
self.polys.append(self.create_poly(self.poly_points, pos=center))
self.poly_points = []
elif event.type == KEYDOWN and event.key == K_SPACE:
self.run_physics = not self.run_physics
elif event.type == KEYDOWN and event.key == K_k:
for x in range (-100,100,25):
for y in range(-100,100,25):
p = pygame.mouse.get_pos()
p = Vec2d(self.flipyv(Vec2d(p))) + (x,y)
self.polys.append(self.create_box(pos=p))
elif event.type == KEYDOWN and event.key == K_b:
p = self.flipyv(Vec2d(pygame.mouse.get_pos()))
self.polys.append(self.create_box(p, size=10, mass = 1))
elif event.type == KEYDOWN and event.key == K_f:
bp = Vec2d(100,500)
p = self.flipyv(Vec2d(pygame.mouse.get_pos())) -bp
ball = self.create_ball(bp)
p = p.normalized()
ball.body.apply_impulse_at_local_point(p*1000, (0,0))
self.balls.append(ball)
elif event.type == KEYDOWN and event.key == K_g:
g = self.space.gravity
g.rotate(45)
self.space.gravity = g
mpos = pygame.mouse.get_pos()
if pygame.key.get_mods() & KMOD_SHIFT and pygame.mouse.get_pressed()[2]:
p = self.flipyv(Vec2d(mpos))
self.poly_points.append(p)
hit = self.space.point_query_nearest( self.flipyv(Vec2d(mpos)), 0, pm.ShapeFilter() )
if hit != None:
self.shape_to_remove = hit.shape
else:
self.shape_to_remove = None
### Update physics
if self.run_physics:
x = 1
dt = 1.0/60.0/x
for x in range(x):
self.space.step(dt)
for ball in self.balls:
#ball.body.reset_forces()
pass
for poly in self.polys:
#poly.body.reset_forces()
pass
### Draw stuff
self.draw()
### Check for objects outside of the screen, we can remove those
# Balls
xs = []
for ball in self.balls:
if ball.body.position.x < -1000 or ball.body.position.x > 1000 \
or ball.body.position.y < -1000 or ball.body.position.y > 1000:
xs.append(ball)
for ball in xs:
self.space.remove(ball, ball.body)
self.balls.remove(ball)
# Polys
xs = []
for poly in self.polys:
if poly.body.position.x < -1000 or poly.body.position.x > 1000 \
or poly.body.position.y < -1000 or poly.body.position.y > 1000:
xs.append(poly)
for poly in xs:
self.space.remove(poly, poly.body)
self.polys.remove(poly)
### Tick clock and update fps in title
self.clock.tick(50)
pygame.display.set_caption("fps: " + str(self.clock.get_fps()))
def main():
demo = PhysicsDemo()
demo.run()
if __name__ == '__main__':
doprof = 0
if not doprof:
main()
else:
import cProfile, pstats
prof = cProfile.run("main()", "profile.prof")
stats = pstats.Stats("profile.prof")
stats.strip_dirs()
stats.sort_stats('cumulative', 'time', 'calls')
stats.print_stats(30)