fix indexing of non-array types

Author: tztsai

src/backend.py

@@ -80,9 +80,10 @@ def _backend(self,level=0):
 
 @extend(node.arrayref)
 def _backend(self,level=0):
-    fmt = "%s[%s]"
-    return fmt % (self.func_expr._backend(),
-                       self.args._backend())
+    from . resolve import A
+    x = self.func_expr._backend()
+    fmt = "%s[%s]" if x in A else "take(%s,%s)"
+    return fmt % (x, self.args._backend())
 
 @extend(node.break_stmt)
 def _backend(self,level=0):
@@ -189,14 +190,10 @@ def _backend(self,level=0):
     if len(self.args) == 1:
         return "%s %s" % (optable.get(self.op,self.op),
                          self.args[0]._backend())
-    if len(self.args) == 2:
-        return "%s %s %s" % (self.args[0]._backend(),
-                           optable.get(self.op,self.op),
-                           self.args[1]._backend())
-    #import pdb;pdb.set_trace()
-    ret = "%s=" % self.ret._backend() if self.ret else ""
-    return ret+"%s(%s)" % (self.op,
-                           ",".join([t._backend() for t in self.args]))
+    if hasattr(self, "ret"):
+        ret = f"{self.ret._backend()}="
+        return ret+"%s(%s)" % (self.op, ",".join([t._backend() for t in self.args]))
+    return optable.get(self.op,self.op).join([t._backend() for t in self.args])
 
 
 @extend(node.expr_list)
@@ -320,8 +317,7 @@ def _backend(self, level=0):
     if not self.args:
         return "[]"
     elif any(b.__class__ is node.string for a in self.args for b in a):
-        s = " + ".join(b._backend() for a in self.args for b in a)
-        return s
+        return " + ".join(b._backend() for a in self.args for b in a)
     else:
         #import pdb; pdb.set_trace()
         return "concat([%s])" % self.args[0]._backend()
@@ -366,12 +362,7 @@ def _backend(self,level=0):
 
 @extend(node.string)
 def _backend(self,level=0):
-    # if '\\' in self.value:
-    #     print(self.value)
-    try:
-        return "'%s'" % str(self.value).encode("string_escape")
-    except:
-        return "'%s'" % str(self.value)
+    return f"'{self.value}'"
 
 @extend(node.sub)
 def _backend(self,level=0):

src/libsmop.py

@@ -17,6 +17,7 @@
 import os
 import sys
 import time
+import json
 import glob
 from sys import stdin, stdout, stderr
 
@@ -28,6 +29,8 @@
 dir = glob.glob
 eps = np.finfo(float).eps
 NaN = np.nan
+jsonencode = json.dumps
+jsondecode = json.loads
 
 
 # def _load_matlab_builtins(*names):
@@ -47,6 +50,18 @@ def clear(*args):
             del globals()[a]
 
 
+def take(a, i):
+    """ Get an item with matlab indexing (1-based). """
+    if isinstance(a, matlabarray):
+        return a[i]
+    else:
+        if isinstance(i, slice):
+            i = slice(i.start - 1, i.stop, i.step)
+        else:
+            i -= 1
+        return a[i]
+
+
 def abs(a):
     return np.abs(a)
 
@@ -76,6 +91,8 @@ def concat(args):
     >>> concat([[1,2,3,4,5] , [1,2,3,4,5]])
     [1, 2, 3, 4, 5, 1, 2, 3, 4, 5]
     """
+    if any(isinstance(a, str) for a in args):
+        return "".join(args)
     t = [matlabarray(a) for a in args]
     return np.concatenate(t)
 
@@ -279,7 +296,7 @@ def length(a):
     try:
         return max(np.asarray(a).shape)
     except ValueError:
-        return 1
+        return len(a)
 
 
 def load(a):

src/resolve.py

@@ -24,6 +24,9 @@
 from . import node
 from . node import extend
 
+
+A = set()  # variables that have array or matrix values
+
 def as_networkx(t):
     G = nx.DiGraph()
     for u in node.postorder(t):
@@ -86,7 +89,7 @@ def _lhs_resolve(self,symtab):
     self.func_expr._resolve(symtab) # A
     self.args._resolve(symtab)      # B
     self.func_expr._lhs_resolve(symtab)
-
+    A.add(self.func_expr.name)
 
 
 @extend(node.expr)
@@ -161,6 +164,9 @@ def _lhs_resolve(self,symtab):
 def _resolve(self,symtab):
     self.args._resolve(symtab)
     self.ret._lhs_resolve(symtab)
+    if isinstance(self.args, (node.matrix, node.cellarray)):
+        if isinstance(self.ret, node.ident):
+            A.add(self.ret.name)
 
 @extend(node.null_stmt)
 @extend(node.continue_stmt)
@@ -210,7 +216,7 @@ def _resolve(self,symtab):
 @extend(node.string)
 @extend(node.comment_stmt)
 def _resolve(self,symtab):
-        pass
+    pass
 
 # @extend(node.call_stmt)
 # def _resolve(self,symtab):
@@ -298,6 +304,8 @@ def fix_let_statement(u):
         if (isinstance(u.ret, node.ident) and
             isinstance(u.args, node.matrix)):
             if any(b.__class__ is node.string for a in u.args.args for b in a):
-                return
-            u.args = node.funcall(func_expr=node.ident("matlabarray"),
-                                  args=node.expr_list([u.args]))
+                u.args = node.expr("+", u.args.args[0])
+                A.discard(u.ret.name)
+            else:
+                u.args = node.funcall(func_expr=node.ident("matlabarray"),
+                                      args=node.expr_list([u.args]))