Commits

Maciej Fijalkowski  committed 4c46bac

implement complex for compile, fix test_zjit

  • Participants
  • Parent commits 95b2087

Comments (0)

Files changed (3)

File pypy/module/micronumpy/compile.py

     def newlist(self, items):
         return ListObject(items)
 
+    def newcomplex(self, r, i):
+        return ComplexObject(r, i)
+
     def listview(self, obj):
         assert isinstance(obj, ListObject)
         return obj.items
             raise OperationError(self.w_TypeError, self.wrap("slice."))
         raise NotImplementedError
 
+    def unpackcomplex(self, w_obj):
+        if isinstance(w_obj, ComplexObject):
+            return w_obj.r, w_obj.i
+        raise NotImplementedError
+
     def index(self, w_obj):
         return self.wrap(self.int_w(w_obj))
 
     def __init__(self, v):
         self.v = v
 
+class ComplexObject(W_Root):
+    tp = FakeSpace.w_complex
+    def __init__(self, r, i):
+        self.r = r
+        self.i = i
+
 class InterpreterState(object):
     def __init__(self, code):
         self.code = code
     def execute(self, interp):
         return interp.space.wrap(self.v)
 
+class ComplexConstant(Node):
+    def __init__(self, r, i):
+        self.r = float(r)
+        self.i = float(i)
+
+    def __repr__(self):
+        return 'ComplexConst(%s, %s)' % (self.r, self.i)
+
+    def wrap(self, space):
+        return space.newcomplex(self.r, self.i)
+
+    def execute(self, interp):
+        return self.wrap(interp.space)
+
 class RangeConstant(Node):
     def __init__(self, v):
         self.v = int(v)
 
     def execute(self, interp):
         w_list = self.wrap(interp.space)
-        dtype = get_dtype_cache(interp.space).w_float64dtype
-        return array(interp.space, w_list, w_dtype=dtype, w_order=None)
+        return array(interp.space, w_list)
 
     def __repr__(self):
         return "[" + ", ".join([repr(item) for item in self.items]) + "]"
                 stack.append(RangeConstant(tokens.pop().v))
                 end = tokens.pop()
                 assert end.name == 'pipe'
+            elif token.name == 'paren_left':
+                stack.append(self.parse_complex_constant(tokens))
             elif accept_comma and token.name == 'comma':
                 continue
             else:
             args += self.parse_expression(tokens, accept_comma=True)
         return FunctionCall(name, args)
 
+    def parse_complex_constant(self, tokens):
+        r = tokens.pop()
+        assert r.name == 'number'
+        assert tokens.pop().name == 'comma'
+        i = tokens.pop()
+        assert i.name == 'number'
+        assert tokens.pop().name == 'paren_right'
+        return ComplexConstant(r.v, i.v)
+
     def parse_array_const(self, tokens):
         elems = []
         while True:
                 elems.append(FloatConstant(token.v))
             elif token.name == 'array_left':
                 elems.append(ArrayConstant(self.parse_array_const(tokens)))
+            elif token.name == 'paren_left':
+                elems.append(self.parse_complex_constant(tokens))
             else:
                 raise BadToken()
             token = tokens.pop()

File pypy/module/micronumpy/test/test_compile.py

         d -> 1
         ''')
         assert interp.results[0].value == 0
+
+    def test_complex(self):
+        interp = self.run('''
+        a = (0, 1)
+        b = [(0, 1), (1, 0)]
+        b -> 0
+        ''')
+        assert interp.results[0].real == 0
+        assert interp.results[0].imag == 1
+        

File pypy/module/micronumpy/test/test_zjit.py

     def setup_class(cls):
         default = """
         a = [1,2,3,4]
+        z = (1, 2)
         c = a + b
         sum(c) -> 1::1
         a -> 3:1:2
                                 'new_with_vtable': 1, 
                                 'int_add': 2, 
                                 'float_ne': 1})
-