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