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             #try:
  76             #    wiiuse.poll(self.wiimotes, self.nmotes)
  77             #except:
  78             #    pass
  79
  80             # allow executing functions in this thread
  81             while True:
  82                 try:
  83                     func, args = self.queue.get_nowait()
  84                 except Empty:
  85                     break
  86                 print 'do:', func.__name__, args
  87                 func(*args)
  88
  89     def do(self, func, *args):
  90         '''Run the function in the thread handling the wiimote'''
  91         self.queue.put((func, args))
  92
  93     def event_cb(self, wmp):
  94         '''Called when the library has some data for the user.'''
  95         wm = wmp[0]
  96         if wm.btns:
  97             for name,b in wiiuse.button.items():
  98                 if wiiuse.is_just_pressed(wm, b):
  99                     pygame.event.post(pygame.event.Event(WIIMOTE_BUTTON_PRESS, button=name,
 100                                                          time=time.time(),
 101                                                          id=wm.unid))
 102
 103         if wm.btns_released:
 104             for name,b in wiiuse.button.items():
 105                 if wiiuse.is_released(wm, b):
 106                     pygame.event.post(pygame.event.Event(WIIMOTE_BUTTON_RELEASE, button=name,
 107                                                          time=time.time(),
 108                                                          id=wm.unid))
 109
 110         if True:
 111             pygame.event.post(pygame.event.Event(WIIMOTE_ACCEL,
 112                                                  orient=(wm.orient.roll, wm.orient.pitch,
 113                                                          wm.orient.yaw),
 114                                                  accel=(wm.gforce.x, wm.gforce.y, wm.gforce.z),
 115                                                  time=time.time(),
 116                                                  id=wm.unid))
 117         if True:
 118             dots = [ (wm.ir.dot[i].visible, wm.ir.dot[i].x, wm.ir.dot[i].y) for i in range(4) ]
 119             pygame.event.post(pygame.event.Event(WIIMOTE_IR,
 120                                                  dots=dots,
 121                                                  cursor=(wm.ir.x, wm.ir.y, wm.ir.z),
 122                                                  time=time.time(),
 123                                                  id=wm.unid))
 124
 125         if wm.exp.type == wiiuse.EXP_NUNCHUK:
 126             nc = wm.exp.u.nunchuk
 127
 128             for name,b in wiiuse.nunchuk_button.items():
 129                 if wiiuse.is_just_pressed(nc, b):
 130                     pygame.event.post(pygame.event.Event(NUNCHUK_BUTTON_PRESS, button=name,
 131                                                          time=time.time(),
 132                                                          id=wm.unid))
 133                 elif wiiuse.is_released(nc, b):
 134                     pygame.event.post(pygame.event.Event(NUNCHUK_BUTTON_RELEASE, button=name,
 135                                                          time=time.time(),
 136                                                          id=wm.unid))
 137
 138             pygame.event.post(pygame.event.Event(NUNCHUK_ACCEL,
 139                                                  orient=(nc.orient.roll, nc.orient.pitch,
 140                                                          nc.orient.yaw),
 141                                                  accel=(nc.gforce.x, nc.gforce.y, nc.gforce.z),
 142                                                  time=time.time(),
 143                                                  id=wm.unid))
 144             pygame.event.post(pygame.event.Event(NUNCHUK_JOY,
 145                                                  angle=nc.js.ang,
 146                                                  mag=nc.js.mag,
 147                                                  time=time.time(),
 148                                                  id=wm.unid))
 149
 150     def control_cb(self, wmp, attachment, speaker, ir, led, battery):
 151         '''Could check the battery level and such here'''
 152         pygame.event.post(pygame.event.Event(WIIMOTE_STATUS,
 153                                              attachment=attachment,
 154                                              speaker=speaker,
 155                                              ir=ir,
 156                                              led=[led[i] for i in range(4)],
 157                                              battery=battery,
 158                                              id=wmp[0].unid))
 159
 160     def disconnect_cb(self, wmp):
 161         '''What should we do here?'''
 162         pygame.event.post(pygame.event.Event(WIIMOTE_DISCONNECT,
 163                                              id=wmp[0].unid))
 164
 165     def quit(self):
 166         '''Go away.'''
 167         for i in range(self.nmotes):
 168             wiiuse.set_leds(self.wiimotes[i], 0)
 169             wiiuse.disconnect(self.wiimotes[i])
 170         self.go = False
 171
 172 WT = None
 173
 174 def init(nmotes, timeout):
 175     '''Initialize the module.'''
 176     global WT
 177     if WT:
 178         return
 179     WT = wiimote_thread(nmotes, timeout)
 180
 181 def get_count():
 182     '''How many Wiimotes were found?'''
 183     return WT.actual_nmotes
 184
 185 def quit():
 186     '''Gracefully shutdown the connection and turn off the wiimote leds'''
 187     WT.quit()
 188     WT.join()
 189
 190 class wiimote(object):
 191     '''Object representing a Wiimote'''
 192     def __init__(self, n):
 193         self.wm = WT.wiimotes[n]
 194
 195     def enable_leds(self, m):
 196         '''Control leds. The lower 4 bits map to the 4 leds'''
 197         WT.do(wiiuse.set_leds, self.wm, sum([wiiuse.LED[i] for i in range(4) if m & (1<<i)]))
 198
 199     def enable_rumble(self, on):
 200         '''Control rumble'''
 201         WT.do(wiiuse.rumble, self.wm, on)
 202
 203     def enable_accels(self, on):
 204         '''Control reporting of accelerometer data.'''
 205         WT.do(wiiuse.motion_sensing, self.wm, on)
 206
 207     def enable_ir(self, on, vres=None, position=None, aspect=None):
 208         '''Control reporting IR data.'''
 209         WT.do(wiiuse.set_ir, self.wm, on)
 210         if vres is not None:
 211             WT.do(wiiuse.set_ir_vres, self.wm, vres)
 212         if position is not None:
 213             WT.do(wiiuse.set_ir_position, self.wm, position)
 214         if aspect is not None:
 215             WT.do(wiiuse.set_aspect_ratio, self.wm, aspect)
 216
 217     def set_flags(self, smoothing=None, continuous=None, threshold=None):
 218         '''Set flags SMOOTHING, CONTINUOUS, ORIENT_THRESH'''
 219         enable = disable = 0
 220         if smoothing is not None:
 221             if smoothing:
 222                 enable |= wiiuse.SMOOTHING
 223             else:
 224                 disable |= wiiuse.SMOOTHING
 225         if continuous is not None:
 226             if continuous:
 227                 enable |= wiiuse.CONTINUOUS
 228             else:
 229                 disable |= wiiuse.CONTINUOUS
 230         if threshold is not None:
 231             if threshold:
 232                 enable |= wiiuse.ORIENT_THRESH
 233             else:
 234                 disable |= wiiuse.ORIENT_THRESH
 235         print enable, disable
 236         WT.do(wiiuse.set_flags, self.wm, enable, disable)
 237
 238     def set_orient_thresh(self, thresh):
 239         '''Set orientation threshold'''
 240         WT.do(wiiuse.set_orient_threshold, self.wm, thresh)
 241
 242     def status(self):
 243         '''Trigger a status callback.'''
 244         WT.do(wiiuse.status, self.wm)
 245
 246     def disconnect(self):
 247         '''Disconnect this Wiimote'''
 248         WT.do(wiiuse.disconnect(self.wm))
 249
 250 def Wiimote(n):
 251     '''Get the object for the nth Wiimote'''
 252     return wiimote(n)
 253