support matching scores as ints (#7)

* support matching scores as ints * Update pytrec_eval.cpp
cvangysel · Feb 17, 2019 · 95b8927 · 95b8927
1 parent f9c04ab
commit 95b8927
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/src/pytrec_eval.cpp b/src/pytrec_eval.cpp
@@ -273,14 +273,14 @@ class ResultRankingBuilder : public RankingBuilder<RESULTS, TEXT_RESULTS_INFO, T
 
     virtual bool ProcessQueryDocumentPair(TEXT_RESULTS* const pair,
                                           PyObject* const inner_value) const {
-        if (!PyFloat_Check(inner_value)) {
-            PyErr_SetString(PyExc_TypeError, "Expected matching score to be float.");
-
-            return false;
+        if (PyFloat_Check(inner_value)) {
+		pair->sim = PyFloat_AsDouble(inner_value);
+	} else if (PyLong_Check(inner_value)) {
+		pair->sim = PyLong_AsDouble(inner_value);
+	} else {
+		PyErr_SetString(PyExc_TypeError, "Expected matching score to be int, long or float.");
+		return false;
         }
-
-        pair->sim = PyFloat_AsDouble(inner_value);
-
         return true;
     }
 };