Commits

Anonymous committed a6be743

24 bit image.tostring, surfarray.pixels* (numpy only)

Comments (0)

Files changed (3)

             surface.lock()
             result = surf.pixels.to_string()
             surface.unlock()
+    elif surf.format.BytesPerPixel == 3 and format in ('RGBA', 'ARGB'):
+        # Optimised conversion from RGB to RGBA or ARGB.
+        if surf.format.Rmask == SDL_SwapLE32(0x000000ff) and \
+           surf.format.Gmask == SDL_SwapLE32(0x0000ff00) and \
+           surf.format.Bmask == SDL_SwapLE32(0x00ff0000) and \
+           pitch == w * surf.format.BytesPerPixel:
+            surface.lock()
+            result = surf.pixels.to_string()
+            surface.unlock()
+
+            # Insert in empty alpha byte
+            alpha = chr(0xff)
+            result = alpha.join(re.findall('...', result, re.DOTALL))
+            if format == 'ARGB':
+                result = alpha + result
+            else:
+                result += alpha
+
     elif surf.format.BytesPerPixel == 4 and format == 'RGB':
         # Optimised conversion from RGBA or ARGB to RGB.
         # This is an optimisation; could also use the default case.
                 pixels = [(palette[c].r, palette[c].g, palette[c].b, 255) \
                           for c in surf.pixels]
         elif surf.format.BytesPerPixel == 3:
-            raise NotImplementedException, 'TODO'
+            raise NotImplementedError, 'TODO'
         else:
             Rmask = surf.format.Rmask
             Gmask = surf.format.Gmask

pygame/surface.py

 __version__ = '$Id$'
 
 from copy import copy
+import weakref
 
 from SDL import *
 import pygame.base
     __slots__ = ['owner', 'pixeloffset', 'offsetx', 'offsety']
 
 class Surface(object):
-    __slots__ = ['_surf', '_subsurface']
+    __slots__ = ['_surf', '_subsurface', '_weakrefs']
 
     def __init__(self, size=(0,0), flags=0, depth=0, masks=None, 
                  surf=None, subsurf=None):
                 Used internally.
 
         '''
+        self._weakrefs = []
         if surf:
             if not isinstance(surf, SDL_Surface):
                 raise TypeError, 'surf'
         '''
         return self._surf._pixels.contents != None
 
+    def lifelock(self, obj):
+        '''Lock the surface for as long as obj is alive.
+
+        This uses a weak reference to obj to detect when it is garbage
+        collected, at which point the surface is unlocked again.
+        '''
+        self.lock()
+        self._weakrefs.append(weakref.ref(obj, self._lifelock_callback))
+
+    def _lifelock_callback(self, ref):
+        self._weakrefs.remove(ref)
+        self.unlock()
+
     def get_at(self, pos):
         '''Get the color value at a single pixel.
 

pygame/surfarray.py

             Surface to reference.
 
     :rtype: Numeric array
-    '''       
+    '''
+    surf = surface._surf
+    bpp = surf.format.BytesPerPixel
+
+    if bpp == 3 or bpp < 1 or bpp > 4:
+        raise ValueError, 'unsupprt bit depth for 2D reference array'
+
+    if surf.pitch % bpp != 0:
+        #TODO hack stride
+        raise NotImplementedError, "Can't correct for this surface pitch"
+
+    shape = surf.h, surf.pitch / bpp
+
+    surface.lock()
+    array = _array_from_buffer(surf.pixels.as_ctypes(), bpp, shape)
+    surface.lifelock(array)
+    surface.unlock()
+
+    array = array[:,:surf.w]
+    return _array.transpose(array)
 
 def array3d(surface):
     '''Copy pixels into a 3d array.
 
     :rtype: Numeric array
     '''  
+    surf = surface._surf
+    bpp = surf.format.BytesPerPixel
+
+    if bpp <= 2 or bpp > 4:
+        raise ValueError, 'unsupport bit depth for alpha array'
+
+    if SDL_SwapLE32(surf.format.Rmask) == 0xff << 16 and \
+       SDL_SwapLE32(surf.format.Gmask) == 0xff << 8 and \
+       SDL_SwapLE32(surf.format.Bmask) == 0xff:
+        start = 2
+        step = -1
+        end = None
+    elif SDL_SwapLE32(surf.format.Rmask) == 0xff and \
+         SDL_SwapLE32(surf.format.Gmask) == 0xff << 8 and \
+         SDL_SwapLE32(surf.format.Bmask) == 0xff << 16:
+        start = 0
+        step = 1
+        end = 3
+    else:
+        raise ValueError, 'unsupport colormasks for 3D reference array'
+
+    shape = surf.h, surf.pitch 
+
+    surface.lock()
+    array = _array_from_buffer(surf.pixels.as_bytes().as_ctypes(), 1, shape)
+    surface.lifelock(array)
+    surface.unlock()
+
+    array = array[:,:surf.w*bpp]
+    array = array.reshape(surf.h, surf.w, bpp)
+    array = array[:,:,start:end:step]
+    return _array.transpose(array, (1, 0, 2))
 
 def array_alpha(surface):
     '''Copy pixel alphas into a 2d array.
     array = array2d(surface)
 
     format = surface._surf.format
-    print format.Amask, format.BytesPerPixel
     if (not format.Amask) or format.BytesPerPixel == 1:
         array[:,:] = 0xff
     else:
-        print 'here'
         array = array >> format.Ashift << format.Aloss
 
     return array.astype(_array.UInt8)
             Surface to reference.
 
     :rtype: Numeric array
-    '''  
+    ''' 
+    surf = surface._surf
+
+    if surf.format.BytesPerPixel != 4:
+        raise ValueError, 'unsupport bit depth for alpha array'
+
+    if SDL_SwapLE32(surf.format.Amask) == 0xff << 24:
+        startpixel = 3
+    else:
+        startpixel = 0
+
+    shape = surf.h, surf.pitch
+
+    surface.lock()
+    array = _array_from_buffer(surf.pixels.as_bytes().as_ctypes(), 1, shape)
+    surface.lifelock(array)
+    surface.unlock()
+
+    array = array[:,:surf.w*4]
+    array = array[:,startpixel::4]
+    return _array.transpose(array)
 
 def array_colorkey(surface):
     '''Copy the colorkey values into a 2d array.
     else:
         import Numeric
         return Numeric
+
+def _array_from_buffer(buffer, bpp, shape):
+    typecode = (_array.UInt8, _array.UInt16, None, _array.UInt32)[bpp-1]
+    if _array.__name__ == 'numpy':
+        return _array.frombuffer(buffer, typecode).reshape(shape)
+    else:
+        raise NotImplementedError, 'numpy required'