Allow constructor to take 3 arguments

moble · Mar 13, 2017 · 53d8322 · 53d8322
1 parent a2c3152
commit 53d8322
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/numpy_quaternion.c b/numpy_quaternion.c
@@ -91,6 +91,8 @@ pyquaternion_init(PyObject *self, PyObject *args, PyObject *kwds)
   // initialization should take place in `tp_new`, while for mutable
   // types, most initialization should be deferred to `tp_init`."
   // ---Python 2.7.8 docs
+
+  Py_ssize_t size = PyTuple_Size(args);
   quaternion* q;
   q = &(((PyQuaternion*)self)->obval);
   if (kwds && PyDict_Size(kwds)) {
@@ -99,10 +101,13 @@ pyquaternion_init(PyObject *self, PyObject *args, PyObject *kwds)
     return -1;
   }
 
-  if (!PyArg_ParseTuple(args, "dddd", &q->w, &q->x, &q->y, &q->z)) {
+  if (((size == 3) && (!PyArg_ParseTuple(args, "ddd", &q->x, &q->y, &q->z)))
+      || ((size == 4) && (!PyArg_ParseTuple(args, "dddd", &q->w, &q->x, &q->y, &q->z)))) {
     PyErr_SetString(PyExc_TypeError,
-                    "quaternion constructor takes four double (float) arguments");
+                    "quaternion constructor takes three or four float arguments");
     return -1;
+  } else if(size == 3) {
+    q->w = 0.0;
   }
 
   return 0;

diff --git a/test/test_quaternion.py b/test/test_quaternion.py
@@ -115,6 +115,13 @@ def test_quaternion_members():
     assert Q.y == 3.3
     assert Q.z == 4.4
 
+    Q = quaternion.quaternion(2.2, 3.3, 4.4)
+    assert Q.real == 0.0
+    assert Q.w == 0.0
+    assert Q.x == 2.2
+    assert Q.y == 3.3
+    assert Q.z == 4.4
+
 
 def test_constants():
     assert quaternion.one == np.quaternion(1.0, 0.0, 0.0, 0.0)