Merge pull request #2284 from Shaikh-Ubaid/syntactic_sugar_for_array

Syntactic sugar for Array annotation

Author: certik

integration_tests/CMakeLists.txt

@@ -424,6 +424,7 @@ RUN(NAME array_01            LABELS cpython llvm wasm c)
 RUN(NAME array_02            LABELS cpython wasm c)
 RUN(NAME array_03            LABELS cpython llvm c)
 RUN(NAME array_04            LABELS cpython llvm c)
+RUN(NAME array_05            LABELS cpython llvm c)
 RUN(NAME bindc_01            LABELS cpython llvm c)
 RUN(NAME bindc_02            LABELS cpython llvm c)
 RUN(NAME bindc_04            LABELS llvm c NOFAST)

integration_tests/array_05.py

@@ -0,0 +1,39 @@
+from lpython import i32, f64, Array
+from numpy import empty, int32, float64
+
+
+def test_1():
+    y: Array[f64, 3] = empty([3], dtype=float64)
+    y[0] = 3.14
+    y[1] = -4.14
+    y[2] = 100.100
+
+    print(y)
+    assert abs(y[0] - (3.14)) <= 1e-6
+    assert abs(y[1] - (-4.14)) <= 1e-6
+    assert abs(y[2] - (100.100)) <= 1e-6
+
+def test_2():
+    x: Array[i32, 2, 3] = empty([2, 3], dtype=int32)
+
+    x[0, 0] = 5
+    x[0, 1] = -10
+    x[0, 2] = 15
+    x[1, 0] = 4
+    x[1, 1] = -14
+    x[1, 2] = 100
+
+    print(x)
+    assert x[0, 0] == 5
+    assert x[0, 1] == -10
+    assert x[0, 2] == 15
+    assert x[1, 0] == 4
+    assert x[1, 1] == -14
+    assert x[1, 2] == 100
+
+
+def main0():
+    test_1()
+    test_2()
+
+main0()

src/lpython/semantics/python_ast_to_asr.cpp

@@ -1851,17 +1851,35 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
                     is_allocatable, raise_error, abi, is_argument);
                 return ASRUtils::TYPE(ASR::make_Const_t(al, loc, type));
             } else {
+                AST::expr_t* dim_info = s->m_slice;
+
+                if (var_annotation == "Array") {
+                    LCOMPILERS_ASSERT(AST::is_a<AST::Tuple_t>(*s->m_slice));
+                    AST::Tuple_t *t = AST::down_cast<AST::Tuple_t>(s->m_slice);
+                    LCOMPILERS_ASSERT(t->n_elts >= 2);
+                    LCOMPILERS_ASSERT(AST::is_a<AST::Name_t>(*t->m_elts[0]));
+                    var_annotation = AST::down_cast<AST::Name_t>(t->m_elts[0])->m_id;
+                    Vec<AST::expr_t*> dims;
+                    dims.reserve(al, 0);
+                    for (size_t i = 1; i < t->n_elts; i++) {
+                        dims.push_back(al, t->m_elts[i]);
+                    }
+                    AST::ast_t* dim_tuple = AST::make_Tuple_t(al, t->base.base.loc, dims.p, dims.size(),
+                        AST::expr_contextType::Load);
+                    dim_info = AST::down_cast<AST::expr_t>(dim_tuple);
+                }
+
                 ASR::ttype_t* type = get_type_from_var_annotation(var_annotation,
                     annotation.base.loc, dims, m_args, n_args, raise_error, abi, is_argument);
 
-                if (AST::is_a<AST::Slice_t>(*s->m_slice)) {
+                if (AST::is_a<AST::Slice_t>(*dim_info)) {
                     ASR::dimension_t dim;
                     dim.loc = loc;
                     dim.m_start = nullptr;
                     dim.m_length = nullptr;
                     dims.push_back(al, dim);
-                } else if( is_runtime_array(s->m_slice) ) {
-                    AST::Tuple_t* tuple_multidim = AST::down_cast<AST::Tuple_t>(s->m_slice);
+                } else if( is_runtime_array(dim_info) ) {
+                    AST::Tuple_t* tuple_multidim = AST::down_cast<AST::Tuple_t>(dim_info);
                     for( size_t i = 0; i < tuple_multidim->n_elts; i++ ) {
                         if( AST::is_a<AST::Slice_t>(*tuple_multidim->m_elts[i]) ) {
                             ASR::dimension_t dim;
@@ -1872,7 +1890,7 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
                         }
                     }
                 } else {
-                    this->visit_expr(*s->m_slice);
+                    this->visit_expr(*dim_info);
                     ASR::expr_t *value = ASRUtils::EXPR(tmp);
                     fill_dims_for_asr_type(dims, value, loc);
                 }

src/runtime/lpython/lpython.py

@@ -77,6 +77,9 @@ def __init__(self, type, dims):
         self._type = type
         self._dims = dims
 
+    def __class_getitem__(self, params):
+        return Array(params[0], params[1:])
+
 i1 = Type("i1")
 i8 = Type("i8")
 i16 = Type("i16")