Distributive-Network · philippedistributive · Jun 14, 2024 · Jun 7, 2024 · Jun 7, 2024 · Jun 7, 2024
diff --git a/src/jsTypeFactory.cc b/src/jsTypeFactory.cc
@@ -294,11 +294,11 @@ JS::Value jsTypeFactorySafe(JSContext *cx, PyObject *object) {
   return v;
 }
 
-void setPyException(JSContext *cx) {
+bool setPyException(JSContext *cx) {
   // Python `exit` and `sys.exit` only raise a SystemExit exception to end the program
   // We definitely don't want to catch it in JS
   if (PyErr_ExceptionMatches(PyExc_SystemExit)) {
-    return;
+    return false;
   }
 
   PyObject *type, *value, *traceback;
@@ -314,52 +314,111 @@ void setPyException(JSContext *cx) {
     JS::RootedValue jsExceptionValue(cx, JS::ObjectValue(*jsException));
     JS_SetPendingException(cx, jsExceptionValue);
   }
+  return true;
 }
 
 bool callPyFunc(JSContext *cx, unsigned int argc, JS::Value *vp) {
   JS::CallArgs callargs = JS::CallArgsFromVp(argc, vp);
 
   // get the python function from the 0th reserved slot
-  JS::Value pyFuncVal = js::GetFunctionNativeReserved(&(callargs.callee()), 0);
-  PyObject *pyFunc = (PyObject *)(pyFuncVal.toPrivate());
-
-  unsigned int callArgsLength = callargs.length();
+  PyObject *pyFunc = (PyObject *)js::GetFunctionNativeReserved(&(callargs.callee()), 0).toPrivate();
+  Py_INCREF(pyFunc);
+  PyObject *pyRval = NULL;
+  PyObject *pyArgs = NULL;
+  Py_ssize_t nNormalArgs = 0;   // number of positional non-default arguments
+  Py_ssize_t nDefaultArgs = 0;  // number of positional default arguments
+  bool varargs = false;
+  bool unknownNargs = false;
+
+  if (PyCFunction_Check(pyFunc)) {
+    const int funcFlags = ((PyCFunctionObject *)pyFunc)->m_ml->ml_flags;
+    if (funcFlags & METH_NOARGS) { // 0 arguments
+      nNormalArgs = 0;
+    }
+    else if (funcFlags & METH_O) { // 1 argument
+      nNormalArgs = 1;
+    }
+    else { // unknown number of arguments
+      nNormalArgs = 0;
+      unknownNargs = true;
+      varargs = true;
+    }
+  }
+  else {
+    nNormalArgs = 1;
+    PyObject *f = pyFunc;
+    if (PyMethod_Check(pyFunc)) {
+      f = PyMethod_Function(pyFunc); // borrowed reference
+      nNormalArgs -= 1; // don't include the implicit `self` of the method as an argument
+    }
+    PyCodeObject *bytecode = (PyCodeObject *)PyFunction_GetCode(f); // borrowed reference
+    PyObject *defaults = PyFunction_GetDefaults(f); // borrowed reference
+    nDefaultArgs = defaults ? PyTuple_Size(defaults) : 0;
+    nNormalArgs += bytecode->co_argcount - nDefaultArgs - 1;
+    if (bytecode->co_flags & CO_VARARGS) {
+      varargs = true;
+    }
+  }
 
-  if (!callArgsLength) {
+  // use faster calling if no arguments are needed
+  if (((nNormalArgs + nDefaultArgs) == 0 && !varargs)) {
     #if PY_VERSION_HEX >= 0x03090000
-    PyObject *pyRval = PyObject_CallNoArgs(pyFunc);
+    pyRval = PyObject_CallNoArgs(pyFunc);
     #else
-    PyObject *pyRval = _PyObject_CallNoArg(pyFunc); // in Python 3.8, the API is only available under the name with a leading underscore
+    pyRval = _PyObject_CallNoArg(pyFunc); // in Python 3.8, the API is only available under the name with a leading underscore
     #endif
-    if (PyErr_Occurred()) { // Check if an exception has already been set in Python error stack
-      setPyException(cx);
-      return false;
+    if (PyErr_Occurred() && setPyException(cx)) { // Check if an exception has already been set in Python error stack
+      goto failure;
     }
-    callargs.rval().set(jsTypeFactory(cx, pyRval));
-    return true;
+    goto success;
   }
 
   // populate python args tuple
-  PyObject *pyArgs = PyTuple_New(callArgsLength);
-  for (size_t i = 0; i < callArgsLength; i++) {
+  Py_ssize_t argTupleLength;
+  if (unknownNargs) { // pass all passed arguments
+    argTupleLength = callargs.length();
+  }
+  else if (varargs) { // if passed arguments is less than number of non-default positionals, rest will be set to `None`
+    argTupleLength = std::max((Py_ssize_t)callargs.length(), nNormalArgs);
+  }
+  else if (nNormalArgs > callargs.length()) { // if passed arguments is less than number of non-default positionals, rest will be set to `None`
+    argTupleLength = nNormalArgs;
+  }
+  else { // passed arguments greater than non-default positionals, so we may be replacing default positional arguments
+    argTupleLength = std::min((Py_ssize_t)callargs.length(), nNormalArgs+nDefaultArgs);
+  }
+  pyArgs = PyTuple_New(argTupleLength);
+
+  for (size_t i = 0; i < callargs.length() && i < argTupleLength; i++) {
     JS::RootedValue jsArg(cx, callargs[i]);
     PyObject *pyArgObj = pyTypeFactory(cx, jsArg);
     if (!pyArgObj) return false; // error occurred
     PyTuple_SetItem(pyArgs, i, pyArgObj);
   }
 
-  PyObject *pyRval = PyObject_Call(pyFunc, pyArgs, NULL);
-  if (PyErr_Occurred()) {
-    setPyException(cx);
-    return false;
+  // set unspecified args to None, to match JS behaviour of setting unspecified args to undefined
+  for (Py_ssize_t i = callargs.length(); i < argTupleLength; i++) {
+    PyTuple_SetItem(pyArgs, i, Py_None);
   }
-  callargs.rval().set(jsTypeFactory(cx, pyRval));
-  if (PyErr_Occurred()) {
-    Py_DECREF(pyRval);
-    setPyException(cx);
-    return false;
+
+  pyRval = PyObject_Call(pyFunc, pyArgs, NULL);
+  if (PyErr_Occurred() && setPyException(cx)) {
+    goto failure;
   }
+  goto success;
 
-  Py_DECREF(pyRval);
+success:
+  if (pyRval) { // can be NULL if SystemExit was raised
+    callargs.rval().set(jsTypeFactory(cx, pyRval));
+    Py_DECREF(pyRval);
+  }
+  Py_DECREF(pyFunc);
+  Py_XDECREF(pyArgs);
   return true;
+
+failure:
+  Py_XDECREF(pyRval);
+  Py_DECREF(pyFunc);
+  Py_XDECREF(pyArgs);
+  return false;
 }
diff --git a/tests/python/test_arrays.py b/tests/python/test_arrays.py
@@ -894,33 +894,6 @@ def myFunc(e, f):
     assert "@evaluate:1:27" in str(e)
 
 
-def test_sort_with_one_arg_keyfunc():
-  items = ['Four', 'Three', 'One']
-
-  def myFunc(e):
-    return len(e)
-  try:
-    pm.eval("(arr, compareFun) => {arr.sort(compareFun)}")(items, myFunc)
-    assert (False)
-  except Exception as e:
-    assert str(type(e)) == "<class 'pythonmonkey.SpiderMonkeyError'>"
-    assert "takes 1 positional argument but 2 were given" in str(e)
-    assert "JS Stack Trace" in str(e)
-    assert "@evaluate:1:27" in str(e)
-
-
-def test_sort_with_builtin_keyfunc():
-  items = ['Four', 'Three', 'One']
-  try:
-    pm.eval("(arr, compareFun) => {arr.sort(compareFun)}")(items, len)
-    assert (False)
-  except Exception as e:
-    assert str(type(e)) == "<class 'pythonmonkey.SpiderMonkeyError'>"
-    assert "len() takes exactly one argument (2 given)" in str(e)
-    assert "JS Stack Trace" in str(e)
-    assert "@evaluate:1:27" in str(e)
-
-
 def test_sort_with_js_func():
   items = ['Four', 'Three', 'One']
   result = [None]

diff --git a/tests/python/test_functions.py b/tests/python/test_functions.py
@@ -0,0 +1,121 @@
+import pythonmonkey as pm
+
+
+def test_func_no_args():
+  def f():
+    return 42
+  assert 42 == pm.eval("(f) => f()")(f)
+
+
+def test_func_too_many_args():
+  def f(a, b):
+    return [a, b]
+  assert [1, 2] == pm.eval("(f) => f(1, 2, 3)")(f)
+
+
+def test_func_equal_args():
+  def f(a, b):
+    return [a, b]
+  assert [1, 2] == pm.eval("(f) => f(1, 2)")(f)
+
+
+def test_func_too_few_args():
+  def f(a, b):
+    return [a, b]
+  assert [1, None] == pm.eval("(f) => f(1)")(f)
+
+
+def test_default_func_no_args():
+  def f(a, b, c=42, d=43):
+    return [a, b, c, d]
+  assert [None, None, 42, 43] == pm.eval("(f) => f()")(f)
+
+
+def test_default_func_too_many_default_args():
+  def f(a, b, c=42, d=43):
+    return [a, b, c, d]
+  assert [1, 2, 3, 4] == pm.eval("(f) => f(1, 2, 3, 4, 5)")(f)
+
+
+def test_default_func_equal_default_args():
+  def f(a, b, c=42, d=43):
+    return [a, b, c, d]
+  assert [1, 2, 3, 4] == pm.eval("(f) => f(1, 2, 3, 4)")(f)
+
+
+def test_default_func_too_few_default_args():
+  def f(a, b, c=42, d=43):
+    return [a, b, c, d]
+  assert [1, 2, 3, 43] == pm.eval("(f) => f(1, 2, 3)")(f)
+
+
+def test_default_func_equal_args():
+  def f(a, b, c=42, d=43):
+    return [a, b, c, d]
+  assert [1, 2, 42, 43] == pm.eval("(f) => f(1, 2)")(f)
+
+
+def test_default_func_too_few_args():
+  def f(a, b, c=42, d=43):
+    return [a, b, c, d]
+  assert [1, None, 42, 43] == pm.eval("(f) => f(1)")(f)
+
+
+def test_vararg_func_no_args():
+  def f(a, b, *args):
+    return [a, b, *args]
+  assert [None, None] == pm.eval("(f) => f()")(f)
+
+
+def test_vararg_func_too_many_args():
+  def f(a, b, *args):
+    return [a, b, *args]
+  assert [1, 2, 3] == pm.eval("(f) => f(1, 2, 3)")(f)
+
+
+def test_vararg_func_equal_args():
+  def f(a, b, *args):
+    return [a, b, *args]
+  assert [1, 2] == pm.eval("(f) => f(1, 2)")(f)
+
+
+def test_vararg_func_too_few_args():
+  def f(a, b, *args):
+    return [a, b, *args]
+  assert [1, None] == pm.eval("(f) => f(1)")(f)
+
+
+def test_default_vararg_func_no_args():
+  def f(a, b, c=42, d=43, *args):
+    return [a, b, c, d, *args]
+  assert [None, None, 42, 43] == pm.eval("(f) => f()")(f)
+
+
+def test_default_vararg_func_too_many_default_args():
+  def f(a, b, c=42, d=43, *args):
+    return [a, b, c, d, *args]
+  assert [1, 2, 3, 4, 5] == pm.eval("(f) => f(1, 2, 3, 4, 5)")(f)
+
+
+def test_default_vararg_func_equal_default_args():
+  def f(a, b, c=42, d=43, *args):
+    return [a, b, c, d, *args]
+  assert [1, 2, 3, 4] == pm.eval("(f) => f(1, 2, 3, 4)")(f)
+
+
+def test_default_vararg_func_too_few_default_args():
+  def f(a, b, c=42, d=43, *args):
+    return [a, b, c, d, *args]
+  assert [1, 2, 3, 43] == pm.eval("(f) => f(1, 2, 3)")(f)
+
+
+def test_default_vararg_func_equal_args():
+  def f(a, b, c=42, d=43, *args):
+    return [a, b, c, d, *args]
+  assert [1, 2, 42, 43] == pm.eval("(f) => f(1, 2)")(f)
+
+
+def test_default_vararg_func_too_few_args():
+  def f(a, b, c=42, d=43, *args):
+    return [a, b, c, d, *args]
+  assert [1, None, 42, 43] == pm.eval("(f) => f(1)")(f)