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