make classes picklable

4 files changed, 61 insertions, 0 deletions

diff --git a/pyisda/curve.pyx b/pyisda/curve.pyx
index bee364c..ef11a42 100644
--- a/pyisda/curve.pyx
+++ b/pyisda/curve.pyx
@@ -1,4 +1,5 @@
 from libc.stdlib cimport malloc, free
+from libc.string cimport memcpy
 from date cimport (JpmcdsStringToDateInterval, pydate_to_TDate, dcc,
                    JpmcdsDateIntervalToFreq, JpmcdsDateFwdThenAdjust, TDate_to_pydate,
                    JpmcdsDateFromBusDaysOffset)
@@ -18,6 +19,24 @@ cdef class Curve:
         if self._thisptr is not NULL:
             JpmcdsFreeTCurve(self._thisptr)
 
+    def __getstate__(self):
+        cdef int num_items = self._thisptr.fNumItems
+        return (num_items,
+                <bytes>(<char*>self._thisptr.fArray)[:sizeof(TRatePt)*num_items],
+                self._thisptr.fBaseDate,
+                self._thisptr.fBasis,
+                self._thisptr.fDayCountConv)
+
+    def __setstate__(self, state):
+        num_items, rates, base_date, basis, dcc = state
+        self._thisptr = <TCurve*>malloc(sizeof(TCurve))
+        self._thisptr.fNumItems = num_items
+        self._thisptr.fArray = <TRatePt*>malloc(sizeof(TRatePt) * num_items)
+        memcpy(self._thisptr.fArray, <char*> rates, sizeof(TRatePt) * num_items)
+        self._thisptr.fBaseDate = base_date
+        self._thisptr.fBasis = basis
+        self._thisptr.fDayCountConv = dcc
+
     def inspect(self):
         """ method to inspect the content of the C struct
 
diff --git a/pyisda/legs.pxd b/pyisda/legs.pxd
index 5c96f62..d711205 100644
--- a/pyisda/legs.pxd
+++ b/pyisda/legs.pxd
@@ -6,6 +6,12 @@ ctypedef int TBoolean
 
 cdef extern from "isda/cx.h":
     ctypedef struct TContingentLeg:
+         TDate         startDate
+         TDate         endDate
+         double        notional
+         TBoolean      protectStart
+
+    ctypedef enum TAccrualPayConv:
         pass
 
     ctypedef struct TFeeLeg:
@@ -13,6 +19,11 @@ cdef extern from "isda/cx.h":
         TDate* accStartDates
         TDate* accEndDates
         TDate* payDates
+        double notional
+        double couponRate
+        long dcc
+        TAccrualPayConv accrualPayConv
+        TBoolean obsStartOfDay
 
 cdef extern from "isda/bastypes.h":
     ctypedef struct TCashFlow:
diff --git a/pyisda/legs.pyx b/pyisda/legs.pyx
index ea19b72..a2d9d0d 100644
--- a/pyisda/legs.pyx
+++ b/pyisda/legs.pyx
@@ -1,5 +1,6 @@
 from libc.stdlib cimport free
 from date cimport pydate_to_TDate, TDate_to_pydate, dcc
+from date import dcc_tostring
 from cdsone cimport JpmcdsStringToStubMethod, TStubMethod
 from curve cimport YieldCurve, SpreadCurve
 
@@ -30,6 +31,12 @@ cdef class ContingentLeg:
         if self._thisptr is not NULL:
             free(self._thisptr)
 
+    def __reduce__(self):
+        return (self.__class__, (TDate_to_pydate(self._thisptr.startDate),
+                                 TDate_to_pydate(self._thisptr.endDate),
+                                 self._thisptr.notional,
+                                 self._thisptr.protectStart))
+
     def pv(self, today, step_in_date, value_date, YieldCurve yc, SpreadCurve sc,
            double recovery_rate):
         """
@@ -100,6 +107,15 @@ cdef class FeeLeg:
         if self._thisptr is NULL:
             raise ValueError
 
+    def __reduce__(self):
+        return (self.__class__, (TDate_to_pydate(self._thisptr.accStartDates[0]),
+                                 TDate_to_pydate(self._thisptr.accEndDates[self._thisptr.nbDates-1]),
+                                 <int>self._thisptr.accrualPayConv,
+                                 self._thisptr.notional,
+                                 self._thisptr.couponRate,
+                                 dcc_tostring(self._thisptr.dcc),
+                                 self._thisptr.obsStartOfDay))
+
     def inspect(self):
         """convenience method to study the C struct"""
         cdef list acc_start_dates = []
diff --git a/tests/test_pickle.py b/tests/test_pickle.py
new file mode 100644
index 0000000..cdb698d
--- /dev/null
+++ b/tests/test_pickle.py
@@ -0,0 +1,15 @@
+import unittest
+from pyisda.legs import ContingentLeg
+from pickle import dumps, loads
+import datetime
+
+class TestPickle(unittest.TestCase):
+
+    def test_contingentleg(self):
+        cl = ContingentLeg(datetime.date(2016, 11, 3),
+                           datetime.date(2021, 12, 20),
+                           12,
+                           True)
+        pickle = dumps(cl)
+        import pdb; pdb.set_trace()
+        cl2 = loads(pickle)