1 import numpy as np
   2 from numpy.testing import *
   3 import obarray
   4 
   5 class TestObarray(NumpyTestCase):
   6     def setUp(self):
   7         
   8         class Foo:
   9             def __init__(self, a, b):
  10                 self.a = a
  11                 self.b = b
  12             def __str__(self):
  13                 return "<Foo a=%s b=%s>" % (self.a, self.b)
  14 
  15         dtype = [("a",np.int),
  16                  ("b",np.float)]
  17         FooArray = obarray.make_obarray(Foo, dtype)
  18 
  19         A = FooArray([Foo(0,0.1),Foo(1,1.2),Foo(2,2.1),Foo(3,3.3)])
  20 
  21         self.Foo = Foo
  22         self.dtype = dtype
  23         self.FooArray = FooArray
  24         self.A = A
  25 
  26     def check_scalar_indexing(self):
  27         f = self.A[0]
  28         self.assert_(isinstance(f,self.Foo))
  29         assert_equal((f.a, f.b), (0,0.1))
  30 
  31     def check_subarray_indexing(self):
  32         A = self.A
  33         assert_equal(type(A[::2]),type(A))
  34         assert_equal(type(A[[True,False,True,True]]),type(A))
  35         assert_equal(type(A[[0,0,0]]),type(A))
  36 
  37     def check_scalar_assignment(self):
  38         Foo = self.Foo
  39         f = Foo(-1,-1.1)
  40         self.A[0] = Foo(-1,-1.1)
  41         assert_equal(self.A[0].a, f.a)
  42         assert_equal(self.A[0].b, f.b)
  43 
  44     def check_vector_assignment(self):
  45         self.A[::2] = self.A[1::2]
  46         assert_equal(self.A[0].a, self.A[1].a)
  47         assert_equal(self.A[0].b, self.A[1].b)
  48 
  49 
  50 if __name__=='__main__':
  51     NumpyTest().run()