Merge pull request #126 from namannimmo10/slice

Implement slice

Author: certik

src/lpython/semantics/python_ast_to_asr.cpp

@@ -823,17 +823,33 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
     void visit_Subscript(const AST::Subscript_t &x) {
         this->visit_expr(*x.m_value);
         ASR::expr_t *value = ASRUtils::EXPR(tmp);
-        this->visit_expr(*x.m_slice);
-        ASR::expr_t *index = ASRUtils::EXPR(tmp);
         Vec<ASR::array_index_t> args;
         args.reserve(al, 1);
         ASR::array_index_t ai;
         ai.loc = x.base.base.loc;
         ai.m_left = nullptr;
-        ai.m_right = index;
+        ai.m_right = nullptr;
         ai.m_step = nullptr;
+        if (AST::is_a<AST::Slice_t>(*x.m_slice)) {
+            AST::Slice_t *s = AST::down_cast<AST::Slice_t>(x.m_slice);
+            if (s->m_lower != nullptr) {
+                this->visit_expr(*s->m_lower);
+                ai.m_left = ASRUtils::EXPR(tmp);
+            }
+            if (s->m_upper != nullptr) {
+                this->visit_expr(*s->m_upper);
+                ai.m_right = ASRUtils::EXPR(tmp);
+            }
+            if (s->m_step != nullptr) {
+                this->visit_expr(*s->m_step);
+                ai.m_step = ASRUtils::EXPR(tmp);
+            }
+        } else {
+            this->visit_expr(*x.m_slice);
+            ASR::expr_t *index = ASRUtils::EXPR(tmp);
+            ai.m_right = index;
+        }
         args.push_back(al, ai);
-
         ASR::symbol_t *s = ASR::down_cast<ASR::Var_t>(value)->m_v;
         ASR::Variable_t *v = ASR::down_cast<ASR::Variable_t>(s);
         ASR::ttype_t *type = v->m_type;

tests/reference/asr-subscript1-1acfc19.json

@@ -0,0 +1,13 @@
+{
+    "basename": "asr-subscript1-1acfc19",
+    "cmd": "lpython --show-asr --no-color {infile} -o {outfile}",
+    "infile": "tests/subscript1.py",
+    "infile_hash": "24efc3980ad13e8772a093253b380c91ce3b38071d69200e5af33d36",
+    "outfile": null,
+    "outfile_hash": null,
+    "stdout": "asr-subscript1-1acfc19.stdout",
+    "stdout_hash": "f7bf6b29c18149b3c536af48eb59a1133156ac03fba83603a083d0fb",
+    "stderr": null,
+    "stderr_hash": null,
+    "returncode": 0
+}

tests/reference/asr-subscript1-1acfc19.stdout

@@ -0,0 +1 @@
+(TranslationUnit (SymbolTable 1 {main_program: (Program (SymbolTable 3 {}) main_program [] []), test_subscript: (Subroutine (SymbolTable 2 {A: (Variable 2 A Local () () Default (Integer 4 [((ConstantInteger 1 (Integer 4 [])) (ConstantInteger 5 (Integer 4 [])))]) Source Public Required .false.), B: (Variable 2 B Local () () Default (Integer 4 [((ConstantInteger 1 (Integer 4 [])) (ConstantInteger 2 (Integer 4 [])))]) Source Public Required .false.), i: (Variable 2 i Local () () Default (Integer 4 []) Source Public Required .false.), s: (Variable 2 s Local () () Default (Character 1 -2 () []) Source Public Required .false.)}) test_subscript [] [(= (Var 2 s) (ConstantString "abc" (Character 1 3 () [])) ()) (= (Var 2 s) (ArrayRef 2 s [(() (ConstantInteger 0 (Integer 4 [])) ())] (Character 1 -2 () []) ()) ()) (= (Var 2 s) (ArrayRef 2 s [((ConstantInteger 1 (Integer 4 [])) (ConstantInteger 2 (Integer 4 [])) ())] (Character 1 -2 () []) ()) ()) (= (Var 2 i) (ArrayRef 2 A [(() (ConstantInteger 0 (Integer 4 [])) ())] (Integer 4 [((ConstantInteger 1 (Integer 4 [])) (ConstantInteger 5 (Integer 4 [])))]) ()) ()) (= (Var 2 B) (ArrayRef 2 A [((ConstantInteger 1 (Integer 4 [])) (ConstantInteger 3 (Integer 4 [])) ())] (Integer 4 [((ConstantInteger 1 (Integer 4 [])) (ConstantInteger 5 (Integer 4 [])))]) ()) ())] Source Public Implementation () .false. .false.)}) [])

tests/reference/ast-subscript1-bd5584b.json

@@ -0,0 +1,13 @@
+{
+    "basename": "ast-subscript1-bd5584b",
+    "cmd": "lpython --show-ast --no-color {infile} -o {outfile}",
+    "infile": "tests/subscript1.py",
+    "infile_hash": "24efc3980ad13e8772a093253b380c91ce3b38071d69200e5af33d36",
+    "outfile": null,
+    "outfile_hash": null,
+    "stdout": "ast-subscript1-bd5584b.stdout",
+    "stdout_hash": "a3078483928581b43987ad69903c0319507833a869b1e1fe7d52edb5",
+    "stderr": null,
+    "stderr_hash": null,
+    "returncode": 0
+}

tests/reference/ast-subscript1-bd5584b.stdout

@@ -0,0 +1 @@
+(Module [(FunctionDef test_subscript ([] [] [] [] [] [] []) [(AnnAssign (Name s Store) (Name str Load) () 1) (Assign [(Name s Store)] (ConstantStr "abc" ()) ()) (Assign [(Name s Store)] (Subscript (Name s Load) (ConstantInt 0 ()) Load) ()) (Assign [(Name s Store)] (Subscript (Name s Load) (Slice (ConstantInt 1 ()) (ConstantInt 2 ()) ()) Load) ()) (AnnAssign (Name A Store) (Subscript (Name i32 Load) (ConstantInt 5 ()) Load) () 1) (AnnAssign (Name B Store) (Subscript (Name i32 Load) (ConstantInt 2 ()) Load) () 1) (AnnAssign (Name i Store) (Name i32 Load) () 1) (Assign [(Name i Store)] (Subscript (Name A Load) (ConstantInt 0 ()) Load) ()) (Assign [(Name B Store)] (Subscript (Name A Load) (Slice (ConstantInt 1 ()) (ConstantInt 3 ()) ()) Load) ())] [] () ())] [])

tests/subscript1.py

@@ -0,0 +1,11 @@
+def test_subscript():
+    s: str
+    s = 'abc'
+    s = s[0]
+    s = s[1:2]
+
+    A: i32[5]
+    B: i32[2]
+    i: i32
+    i = A[0]
+    B = A[1:3]

tests/tests.toml

@@ -23,6 +23,11 @@ filename = "constants1.py"
 ast = true
 asr = true
 
+[[test]]
+filename = "subscript1.py"
+ast = true
+asr = true
+
 [[test]]
 filename = "list1.py"
 ast = true