make it python2 compatible

4 files changed, 18 insertions, 11 deletions

diff --git a/pyisda/curve.pyx b/pyisda/curve.pyx
index e2e602e..82d214b 100644
--- a/pyisda/curve.pyx
+++ b/pyisda/curve.pyx
@@ -14,7 +14,7 @@ cpdef public enum BadDay:
     NONE = <long>'N'
     MODIFIED = <long>'M'
 
-cdef class Curve:
+cdef class Curve(object):
     def __dealloc__(self):
         if self._thisptr is not NULL:
             JpmcdsFreeTCurve(self._thisptr)
diff --git a/pyisda/date.pyx b/pyisda/date.pyx
index bc6a05e..3cdb394 100644
--- a/pyisda/date.pyx
+++ b/pyisda/date.pyx
@@ -1,5 +1,6 @@
 import datetime
 from cpython cimport datetime as c_datetime
+from cpython.version cimport PY_MAJOR_VERSION
 
 c_datetime.import_datetime()
 
@@ -23,5 +24,8 @@ cdef long dcc(str day_count) except -1:
 
 def dcc_tostring(long day_count):
     cdef char* c_string = JpmcdsFormatDayCountConv(day_count)
-    s = c_string.decode('utf-8')
+    if PY_MAJOR_VERSION >= 3:
+        s = c_string.decode('utf-8')
+    else:
+        s = c_string
     return s
diff --git a/pyisda/utils.py b/pyisda/utils.py
index b6724ce..8fda278 100644
--- a/pyisda/utils.py
+++ b/pyisda/utils.py
@@ -7,7 +7,7 @@ from quantlib.currency.api import USDCurrency, EURCurrency
 from quantlib.indexes.ibor_index import IborIndex
 from quantlib.termstructures.yields.api import (
     PiecewiseYieldCurve, DepositRateHelper, SwapRateHelper)
-import array
+import numpy as np
 import datetime
 import requests
 import zipfile
@@ -77,7 +77,7 @@ def build_yc(trade_date, ql_curve = False):
         settings.evaluation_date = Date.from_datetime(trade_date)
         yield_helpers = rate_helpers(MarkitData = markit_data)
         ql_yc = YC(helpers = yield_helpers)
-        dfs = array.array('d', [ql_yc.discount(yh.latest_date) for yh in yield_helpers])
+        dfs = np.array([ql_yc.discount(yh.latest_date) for yh in yield_helpers])
         dates = [pydate_from_qldate(yh.latest_date) for yh in yield_helpers]
         yc = YieldCurve.from_discount_factors(trade_date, dates, dfs, 'ACT/365F')
     else:
@@ -86,7 +86,7 @@ def build_yc(trade_date, ql_curve = False):
         rates = list(rates)
         periods_swaps, rates_swaps = zip(*markit_data['swaps'])
         types = 'M' * len(periods) + 'S' * len(periods_swaps)
-        rates = array.array('d', rates)
+        rates = np.array(rates)
         periods += periods_swaps
         rates += rates_swaps
         yc = YieldCurve(trade_date, types, periods, rates, 'ACT/360', '6M',
diff --git a/tests/test_pickle.py b/tests/test_pickle.py
index 4bb1d46..6417c86 100644
--- a/tests/test_pickle.py
+++ b/tests/test_pickle.py
@@ -1,5 +1,5 @@
 import unittest
-import array
+import numpy as np
 from pyisda.legs import ContingentLeg, FeeLeg
 from pyisda.curve import SpreadCurve
 from pickle import dumps, loads
@@ -18,7 +18,7 @@ class TestPickle(unittest.TestCase):
 
     def test_yc(self):
         orig_dfs = [self.yc.discount_factor(d) for d in self.yc.dates]
-        pickled_yc = loads(dumps(self.yc))
+        pickled_yc = loads(dumps(self.yc, 2))
         self.assertListAlmostEqual([pickled_yc.discount_factor(d) for d in pickled_yc.dates],
                                    orig_dfs)
 
@@ -30,13 +30,16 @@ class TestPickle(unittest.TestCase):
         cl = ContingentLeg(start_date, end_date, 1)
         fl = FeeLeg(start_date, end_date, True, 1, 1)
         sc = SpreadCurve(self.trade_date, self.yc, start_date,
-                         step_in_date, value_date, [end_date], array.array('d', [75*1e-4]), 0.4)
-        sc_pickled = loads(dumps(sc))
+                         step_in_date, value_date, [end_date], np.array([75*1e-4]), 0.4)
+        sc_pickled = loads(dumps(sc, 2))
         self.assertListAlmostEqual([sc.survival_probability(d) for d in self.yc.dates],
                                    [sc_pickled.survival_probability(d) for d in self.yc.dates])
-        cl_pickled = loads(dumps(cl))
-        fl_pickled = loads(dumps(fl))
+        cl_pickled = loads(dumps(cl, 2))
+        fl_pickled = loads(dumps(fl, 2))
         self.assertEqual(cl_pickled.pv(self.trade_date, step_in_date, value_date, self.yc, sc, 0.4),
                          cl.pv(self.trade_date, step_in_date, value_date, self.yc, sc, 0.4))
         self.assertEqual(fl.pv(self.trade_date, step_in_date, value_date, self.yc, sc, False),
                          fl_pickled.pv(self.trade_date, step_in_date, value_date, self.yc, sc, False))
+
+if __name__=="__main__":
+    unittest.main()