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