Un peu de refactoring.
[minwii.git] / src / pywiiuse / pygame_wiimote.py
1 '''Access to the Wiimote in pygame
2
3 I tried to mimic the Joystick interface already in pygame.
4
5 The usuage pattern is
6
7 init the module telling it the maximum number of wiimotes and the timeout
8 get_count to determine how many you got
9 Wiimote(n) to get an object referencing the nth wiimote
10
11 Free for any use. If you or your lawyer are stupid enough to believe I have any liability for
12 this, then don't use it; otherwise be my guest.
13
14 Gary Bishop January 2008
15
16 '''
17
18 import pygame
19 from threading import Thread
20 from Queue import Queue, Empty
21 import time
22
23 # events to use. Is there a way to get ones known to be unused?
24 base = pygame.USEREVENT
25 WIIMOTE_BUTTON_PRESS = base + 1
26 WIIMOTE_BUTTON_RELEASE = base + 2
27 WIIMOTE_ACCEL = base + 3
28 WIIMOTE_IR = base + 4
29 NUNCHUK_BUTTON_PRESS = base + 5
30 NUNCHUK_BUTTON_RELEASE = base + 6
31 NUNCHUK_ACCEL = base + 7
32 NUNCHUK_JOY = base + 8
33 WIIMOTE_STATUS = base + 9
34 WIIMOTE_DISCONNECT = base + 10
35
36 wiiuse = None # import within the thread, why do I have to do this?
37
38 class wiimote_thread(Thread):
39 '''Manage the wiiuse interface'''
40 def __init__(self, nmotes=1, timeout=5):
41 Thread.__init__(self, name='wiimote')
42 self.queue = Queue()
43 self.startup = Queue()
44 self.nmotes = nmotes
45 self.timeout = timeout
46 self.setDaemon(1)
47 self.start()
48 self.startup.get(True) # wait for the thread to get started and acquire the motes
49
50 def run(self):
51 '''This runs in a separate thread'''
52 global wiiuse
53 import PyWiiUse as wiiuse # import here to avoid thread problems on windows
54 self.wiimotes = wiiuse.init(self.nmotes)
55 found = wiiuse.find(self.wiimotes, self.nmotes, self.timeout)
56 self.actual_nmotes = wiiuse.connect(self.wiimotes, self.nmotes)
57
58 for i in range(self.nmotes):
59 wiiuse.set_leds(self.wiimotes[i], wiiuse.LED[i])
60
61 self.go = self.actual_nmotes != 0
62
63 self.startup.put(self.go)
64
65 while self.go:
66 try :
67 if wiiuse.poll(self.wiimotes, self.nmotes) :
68 for i in range(self.nmotes):
69 m = self.wiimotes[i]
70 if m[0].event == wiiuse.EVENT:
71 self.event_cb(m)
72 except :
73 pass
74
75 # allow executing functions in this thread
76 while True:
77 try:
78 func, args = self.queue.get_nowait()
79 except Empty:
80 break
81 func(*args)
82
83 def do(self, func, *args):
84 '''Run the function in the thread handling the wiimote'''
85 self.queue.put((func, args))
86
87 def event_cb(self, wmp):
88 '''Called when the library has some data for the user.'''
89 wm = wmp[0]
90 if wm.btns:
91 for name,b in wiiuse.button.items():
92 if wiiuse.is_just_pressed(wm, b):
93 pygame.event.post(pygame.event.Event(WIIMOTE_BUTTON_PRESS, button=name,
94 time=time.time(),
95 id=wm.unid))
96
97 if wm.btns_released:
98 for name,b in wiiuse.button.items():
99 if wiiuse.is_released(wm, b):
100 pygame.event.post(pygame.event.Event(WIIMOTE_BUTTON_RELEASE, button=name,
101 time=time.time(),
102 id=wm.unid))
103
104 if True:
105 pygame.event.post(pygame.event.Event(WIIMOTE_ACCEL,
106 orient=(wm.orient.roll, wm.orient.pitch,
107 wm.orient.yaw),
108 accel=(wm.gforce.x, wm.gforce.y, wm.gforce.z),
109 time=time.time(),
110 id=wm.unid))
111 if True:
112 dots = [ (wm.ir.dot[i].visible, wm.ir.dot[i].x, wm.ir.dot[i].y) for i in range(4) ]
113 pygame.event.post(pygame.event.Event(WIIMOTE_IR,
114 dots=dots,
115 cursor=(wm.ir.x, wm.ir.y, wm.ir.z),
116 time=time.time(),
117 id=wm.unid))
118
119 if wm.exp.type == wiiuse.EXP_NUNCHUK:
120 nc = wm.exp.u.nunchuk
121
122 for name,b in wiiuse.nunchuk_button.items():
123 if wiiuse.is_just_pressed(nc, b):
124 pygame.event.post(pygame.event.Event(NUNCHUK_BUTTON_PRESS, button=name,
125 time=time.time(),
126 id=wm.unid))
127 elif wiiuse.is_released(nc, b):
128 pygame.event.post(pygame.event.Event(NUNCHUK_BUTTON_RELEASE, button=name,
129 time=time.time(),
130 id=wm.unid))
131
132 pygame.event.post(pygame.event.Event(NUNCHUK_ACCEL,
133 orient=(nc.orient.roll, nc.orient.pitch,
134 nc.orient.yaw),
135 accel=(nc.gforce.x, nc.gforce.y, nc.gforce.z),
136 time=time.time(),
137 id=wm.unid))
138 pygame.event.post(pygame.event.Event(NUNCHUK_JOY,
139 angle=nc.js.ang,
140 mag=nc.js.mag,
141 time=time.time(),
142 id=wm.unid))
143
144 def control_cb(self, wmp, attachment, speaker, ir, led, battery):
145 '''Could check the battery level and such here'''
146 pygame.event.post(pygame.event.Event(WIIMOTE_STATUS,
147 attachment=attachment,
148 speaker=speaker,
149 ir=ir,
150 led=[led[i] for i in range(4)],
151 battery=battery,
152 id=wmp[0].unid))
153
154 def disconnect_cb(self, wmp):
155 '''What should we do here?'''
156 pygame.event.post(pygame.event.Event(WIIMOTE_DISCONNECT,
157 id=wmp[0].unid))
158
159 def quit(self):
160 '''Go away.'''
161 for i in range(self.nmotes):
162 wiiuse.set_leds(self.wiimotes[i], 0)
163 wiiuse.disconnect(self.wiimotes[i])
164 self.go = False
165
166 WT = None
167
168 def init(nmotes, timeout):
169 '''Initialize the module.'''
170 global WT
171 if WT:
172 return
173 WT = wiimote_thread(nmotes, timeout)
174
175 def get_count():
176 '''How many Wiimotes were found?'''
177 return WT.actual_nmotes
178
179 def quit():
180 '''Gracefully shutdown the connection and turn off the wiimote leds'''
181 WT.quit()
182 WT.join()
183
184 class wiimote(object):
185 '''Object representing a Wiimote'''
186 def __init__(self, n):
187 self.wm = WT.wiimotes[n]
188
189 def enable_leds(self, m):
190 '''Control leds. The lower 4 bits map to the 4 leds'''
191 WT.do(wiiuse.set_leds, self.wm, sum([wiiuse.LED[i] for i in range(4) if m & (1<<i)]))
192
193 def enable_rumble(self, on):
194 '''Control rumble'''
195 WT.do(wiiuse.rumble, self.wm, on)
196
197 def enable_accels(self, on):
198 '''Control reporting of accelerometer data.'''
199 WT.do(wiiuse.motion_sensing, self.wm, on)
200
201 def enable_ir(self, on, vres=None, position=None, aspect=None):
202 '''Control reporting IR data.'''
203 WT.do(wiiuse.set_ir, self.wm, on)
204 if vres is not None:
205 WT.do(wiiuse.set_ir_vres, self.wm, vres)
206 if position is not None:
207 WT.do(wiiuse.set_ir_position, self.wm, position)
208 if aspect is not None:
209 WT.do(wiiuse.set_aspect_ratio, self.wm, aspect)
210
211 def set_flags(self, smoothing=None, continuous=None, threshold=None):
212 '''Set flags SMOOTHING, CONTINUOUS, ORIENT_THRESH'''
213 enable = disable = 0
214 if smoothing is not None:
215 if smoothing:
216 enable |= wiiuse.SMOOTHING
217 else:
218 disable |= wiiuse.SMOOTHING
219 if continuous is not None:
220 if continuous:
221 enable |= wiiuse.CONTINUOUS
222 else:
223 disable |= wiiuse.CONTINUOUS
224 if threshold is not None:
225 if threshold:
226 enable |= wiiuse.ORIENT_THRESH
227 else:
228 disable |= wiiuse.ORIENT_THRESH
229 print enable, disable
230 WT.do(wiiuse.set_flags, self.wm, enable, disable)
231
232 def set_orient_thresh(self, thresh):
233 '''Set orientation threshold'''
234 WT.do(wiiuse.set_orient_threshold, self.wm, thresh)
235
236 def status(self):
237 '''Trigger a status callback.'''
238 WT.do(wiiuse.status, self.wm)
239
240 def disconnect(self):
241 '''Disconnect this Wiimote'''
242 WT.do(wiiuse.disconnect(self.wm))
243
244 def Wiimote(n):
245 '''Get the object for the nth Wiimote'''
246 return wiimote(n)
247