Merge pull request #1780 from Shaikh-Ubaid/numpy_size_test_case

Fix numpy size() and add test

Author: certik

integration_tests/CMakeLists.txt

@@ -281,6 +281,7 @@ RUN(NAME variable_decl_03         LABELS cpython llvm c)
 RUN(NAME array_expr_01            LABELS cpython llvm c)
 RUN(NAME array_expr_02            LABELS cpython llvm c)
 RUN(NAME array_size_01            LABELS cpython llvm c)
+RUN(NAME array_size_02            LABELS cpython llvm c)
 RUN(NAME array_01            LABELS cpython llvm wasm c)
 RUN(NAME array_02            LABELS cpython wasm c)
 RUN(NAME bindc_01            LABELS cpython llvm c)

integration_tests/array_size_02.py

@@ -0,0 +1,90 @@
+from lpython import i32, f64, c32, c64, u32
+from numpy import empty, size
+
+def main0():
+    x: i32[4, 5, 2] = empty([4, 5, 2])
+    y: f64[24, 100, 2, 5] = empty([24, 100, 2, 5])
+    z: i32
+    w: i32
+    z = 2
+    w = 3
+    print(size(x))
+    print(size(x, 0))
+    print(size(x, 1))
+    print(size(x, 2))
+    print(size(y))
+    print(size(y, 0))
+    print(size(y, 1))
+    print(size(y, z))
+    print(size(y, w))
+
+    assert size(x) == 40
+    assert size(x, 0) == 4
+    assert size(x, 1) == 5
+    assert size(x, 2) == 2
+    assert size(y) == 24000
+    assert size(y, 0) == 24
+    assert size(y, 1) == 100
+    assert size(y, z) == 2
+    assert size(y, w) == 5
+
+def main1():
+    a: c32[12] = empty([12])
+    b: c64[15, 15, 10] = empty([15, 15, 10])
+    c: i32
+    d: i32
+    c = 1
+    d = 2
+    print(size(a))
+    print(size(a, 0))
+    print(size(b))
+    print(size(b, 0))
+    print(size(b, c))
+    print(size(b, d))
+
+    assert size(a) == 12
+    assert size(a, 0) == 12
+    assert size(b) == 2250
+    assert size(b, 0) == 15
+    assert size(b, c) == 15
+    assert size(b, d) == 10
+
+def main2():
+    a: i32[2, 3] = empty([2, 3])
+    print(size(a))
+    print(size(a, 0))
+    print(size(a, 1))
+
+    assert size(a) == 2*3
+    assert size(a, 0) == 2
+    assert size(a, 1) == 3
+
+def main3():
+    a: u32[2, 3, 4] = empty([2, 3, 4])
+    b: u64[10, 5] = empty([10, 5])
+    c: i32
+    d: i32
+    c = 1
+    d = 2
+    print(size(a))
+    print(size(a, 0))
+    print(size(a, c))
+    print(size(a, d))
+
+    print(size(b))
+    print(size(b, 0))
+    print(size(b, c))
+
+    assert size(a) == 2*3*4
+    assert size(a, 0) == 2
+    assert size(a, c) == 3
+    assert size(a, d) == 4
+
+    assert size(b) == 50
+    assert size(b, 0) == 10
+    assert size(b, c) == 5
+
+main0()
+main1()
+main2()
+main3()

src/lpython/semantics/python_ast_to_asr.cpp

@@ -6749,13 +6749,16 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
                                         std::to_string(args.size()) + " arguments instead.",
                                         x.base.base.loc);
                 }
+                const Location &loc = x.base.base.loc;
                 ASR::expr_t *var = args[0].m_value;
                 ASR::expr_t *dim = nullptr;
+                ASR::ttype_t *int_type = ASRUtils::TYPE(ASR::make_Integer_t(al, loc, 4, nullptr, 0));
                 if (args.size() == 2) {
-                    dim = args[1].m_value;
+                    ASR::expr_t* const_one = ASRUtils::EXPR(make_IntegerConstant_t(al, loc, 1, int_type));
+                    dim = ASRUtils::EXPR(ASR::make_IntegerBinOp_t(al, loc,
+                        args[1].m_value, ASR::binopType::Add, const_one, int_type, nullptr));
                 }
-                ASR::ttype_t *int_type = ASRUtils::TYPE(ASR::make_Integer_t(al, x.base.base.loc, 4, nullptr, 0));
-                tmp = ASR::make_ArraySize_t(al, x.base.base.loc, var, dim, int_type, nullptr);
+                tmp = ASR::make_ArraySize_t(al, loc, var, dim, int_type, nullptr);
                 return;
             } else if (call_name == "empty") {
                 // TODO: check that the `empty` arguments are compatible