feat(augassign): impl a.__iXX__(b) fallback to a=a.__XX__(b)

Author: litlighilit

Objects/pyobject.nim

@@ -21,11 +21,13 @@ export pyobjectBase
 template getMagic*(obj: PyObject, methodName): untyped = 
   obj.pyType.magicMethods.methodName
 
-template getFun*(obj: PyObject, methodName: untyped, handleExcp=false): untyped = 
-  if obj.pyType.isNil:
-    unreachable("Py type not set")
-  let fun = getMagic(obj, methodName)
-  if fun.isNil:
+
+template checkTypeNotNil(obj) =
+  when not defined(release):
+    if obj.pyType.isNil:
+      unreachable("Py type not set")
+
+template handleNilFunOfGetFun(obj, methodName, handleExcp) =
     let objTypeStr = $obj.pyType.name
     let methodStr = astToStr(methodName)
     let msg = "No " & methodStr & " method for " & objTypeStr & " defined"
@@ -34,33 +36,56 @@ template getFun*(obj: PyObject, methodName: untyped, handleExcp=false): untyped
       handleException(excp)
     else:
       return excp
+
+template getFun*(obj: PyObject, methodName: untyped, handleExcp=false): untyped = 
+  bind checkTypeNotNil, handleNilFunOfGetFun
+  obj.checkTypeNotNil
+  let fun = getMagic(obj, methodName)
+  if fun.isNil:
+    handleNilFunOfGetFun(obj, methodName, handleExcp)
   fun
 
 
 # XXX: `obj` is used twice so it better be a simple identity
 # if it's a function then the function is called twice!
 
-# is there any ways to reduce the repetition? simple template won't work
-template callMagic*(obj: PyObject, methodName: untyped, handleExcp=false): PyObject = 
-  let fun = obj.getFun(methodName, handleExcp)
-  let res = fun(obj)
+template checkExcAndRet[T](res: T, handleExcp): T =
   when handleExcp:
     if res.isThrownException:
       handleException(res)
     res
   else:
     res
 
+# is there any ways to reduce the repetition? simple template won't work
+template callMagic*(obj: PyObject, methodName: untyped, handleExcp=false): PyObject = 
+  let fun = obj.getFun(methodName, handleExcp)
+  let res = fun(obj)
+  bind checkExcAndRet
+  res.checkExcAndRet handleExcp
+
 
 template callMagic*(obj: PyObject, methodName: untyped, arg1: PyObject, handleExcp=false): PyObject = 
   let fun = obj.getFun(methodName, handleExcp)
   let res = fun(obj, arg1)
-  when handleExcp:
-    if res.isThrownException:
-      handleException(res)
-    res
+  bind checkExcAndRet
+  res.checkExcAndRet handleExcp
+
+
+template callInplaceMagic*(obj: PyObject, methodName1: untyped,
+    arg1: PyObject, handleExcp=false): PyObject = 
+  bind checkTypeNotNil, handleNilFunOfGetFun
+  bind checkExcAndRet
+  obj.checkTypeNotNil
+  var fun = obj.getMagic(methodName1)
+  if fun.isNil:
+    pyNotImplemented
   else:
-    res
+    if fun.isNil:
+      handleNilFunOfGetFun(obj, methodName1, handleExcp)
+    else:
+      let res = fun(obj, arg1)
+      res.checkExcAndRet handleExcp
 
 template callMagic*(obj: PyObject, methodName: untyped, 
                     arg1, arg2: PyObject, handleExcp=false): PyObject = 

Python/neval.nim

@@ -10,7 +10,7 @@ import builtindict
 import traceback
 import ../Objects/[pyobject, baseBundle, tupleobject, listobject, dictobject,
                    sliceobject, codeobject, frameobject, funcobject, cellobject,
-                   setobject,
+                   setobject, notimplementedobject,
                    exceptionsImpl, moduleobject, methodobject]
 import ../Utils/utils
 
@@ -46,17 +46,46 @@ template doUnary(opName: untyped) =
   let res = top.callMagic(opName, handleExcp=true)
   sSetTop res
 
-macro callInplaceMagic(op1, instr, op2): untyped =
-  let iMagic = ident 'i' & instr.strVal
+macro tryCallInplaceMagic(op1, opName, op2): untyped =
+  let iMagic = ident 'i' & opName.strVal
   quote do:
-    `op1`.callMagic(`iMagic`, `op2`, handleExcp=true)
+    `op1`.callInplaceMagic(`iMagic`, `op2`, handleExcp=false)
 
 template doInplace(opName: untyped) =
   bind callInplaceMagic
   let op2 = sPop()
   let op1 = sTop()
-  let res = op1.callInplaceMagic(opName, op2)
-  sSetTop res
+  let res = op1.tryCallInplaceMagic(opName, op2)
+  if res.isNotImplemented:
+    var nres = op1.callMagic(opName, op2, handleExcp=true)
+    if nres.isNotImplemented:
+      let
+        opStr{.inject.} = "i" & astToStr(opName)
+        # PY-DIFF: not iadd, but +=
+        typ1{.inject.} = op1.pyType.name
+        typ2{.inject.} = op2.pyType.name
+      nres = newTypeError(
+        &"unsupported operand type(s) for '{opStr}': '{typ1}' and '{typ2}'"
+      )
+      sSetTop nres
+    else:
+      sSetTop nres
+      let stIdx = lastI - 2
+      let (opCode, opArg) = f.code.code[stIdx]
+      var st: OpCode
+      case opCode
+      of OpCode.LoadFast: st=StoreFast
+      of OpCode.LoadGlobal: st=StoreGlobal
+      of OpCode.LoadAttr: st=StoreAttr
+      of OpCode.LoadName: st=StoreName
+      of OpCode.LoadDeref: st=StoreDeref
+      of {OpCode.LoadClosure, LoadMethod, LoadClassDeref}:
+        raiseAssert(
+          &"augumented assignment for {opCode} is not implemented yet")
+      else: unreachable()
+      f.code.code.insert (st, opArg), lastI+1
+  else:
+    sSetTop res
 
 template doBinary(opName: untyped) =
   let op2 = sPop()