add an enum for the buffer formats

2 files changed, 20 insertions, 14 deletions

diff --git a/pyisda/curve.pxd b/pyisda/curve.pxd
index d0d85f6..f12fb54 100644
--- a/pyisda/curve.pxd
+++ b/pyisda/curve.pxd
@@ -211,7 +211,7 @@ cdef class Curve:
     cdef readonly size_t buf_size
     cdef inline const TCurve* get_TCurve(self) nogil
 
-    cpdef bytes as_bytes(self, int fmt)
+    cpdef bytes as_bytes(self, Fmt fmt)
 
 cdef class YieldCurve(Curve):
     pass
@@ -232,3 +232,8 @@ cdef void tweak_curve(const TCurve* sc, TCurve* sc_tweaked, double epsilon,
 cdef uint16_t name_offset_from_buf(const char* buf) noexcept nogil
 
 cdef size_t buf_size(int n, int ticker_len) noexcept nogil
+
+cpdef enum Fmt:
+    Plain = 0
+    LZ4 = 1
+    Packed = 2
diff --git a/pyisda/curve.pyx b/pyisda/curve.pyx
index 1322c9e..5ae18b8 100644
--- a/pyisda/curve.pyx
+++ b/pyisda/curve.pyx
@@ -2,7 +2,6 @@
 from cython.operator import dereference as deref, preincrement as preinc
 from cpython.bytes cimport PyBytes_FromStringAndSize, PyBytes_AS_STRING
 from libc.math cimport log1p, log, exp, isnan
-from libc.stdint cimport uint16_t
 from libc.string cimport strcpy, strncpy, strlen
 from .date cimport (JpmcdsStringToDateInterval, pydate_to_TDate, dcc,
                    JpmcdsDateIntervalToFreq, JpmcdsDateFwdThenAdjust, TDate_to_pydate,
@@ -71,7 +70,7 @@ cdef class Curve(object):
     cdef inline const TCurve* get_TCurve(self) nogil:
         return <TCurve*>self.buf.get()
 
-    cpdef bytes as_bytes(self, int fmt):
+    cpdef bytes as_bytes(self, Fmt fmt):
         """ 0 uncompressed
             1 lz4 compressed
             2 packed
@@ -87,15 +86,15 @@ cdef class Curve(object):
             const TCurve* curve = self.get_TCurve()
             TDate ref_date
 
-        if fmt == 1:
+        if fmt == Fmt.LZ4:
             dst_capacity = LZ4_compressBound(size)
             dst = <char*>malloc(dst_capacity)
             compressed_size = LZ4_compress_default(self.buf.get(), dst, size, dst_capacity)
             r = dst[:compressed_size]
             free(dst)
-        elif fmt == 0:
+        elif fmt == Fmt.Plain:
             r = self.buf.get()[:self.buf_size]
-        elif fmt == 2:
+        elif fmt == Fmt.Packed:
             packed_size = size - curve.fNumItems * (sizeof(TDate) - sizeof(uint16_t))
             r = PyBytes_FromStringAndSize(NULL, packed_size)
             dst = PyBytes_AS_STRING(r)
@@ -116,7 +115,7 @@ cdef class Curve(object):
         return r
 
     def __getstate__(self):
-        return self.as_bytes(0)
+        return self.as_bytes(Fmt.Plain)
 
     def __setstate__(self, bytes state not None):
         cdef:
@@ -131,7 +130,7 @@ cdef class Curve(object):
             self.buf.reset(curve, char_free)
 
     @classmethod
-    def from_bytes(cls, object state, int fmt=1):
+    def from_bytes(cls, object state, Fmt fmt=Fmt.LZ4):
         cdef:
             Curve instance = cls.__new__(cls)
             Py_buffer* py_buf
@@ -153,7 +152,7 @@ cdef class Curve(object):
             size = PyBytes_GET_SIZE(state)
 
         with nogil:
-            if fmt == 1:
+            if fmt == Fmt.LZ4:
                 while True:
                     curve = <char*>realloc(curve, decomp_size)
                     state_size = LZ4_decompress_safe(src, curve, size, decomp_size)
@@ -166,11 +165,11 @@ cdef class Curve(object):
                             decomp_size *= 2
                     else:
                         break
-            elif fmt == 0:
+            elif fmt == Fmt.Plain:
                 curve = <char*>malloc(size)
                 memcpy(curve, src, size)
                 state_size = size
-            elif fmt == 2:
+            elif fmt == Fmt.Packed:
                 n = (<TCurve*>src).fNumItems
                 state_size = size + n * (sizeof(TDate) - sizeof(uint16_t))
                 curve = <char*>malloc(state_size)
@@ -186,6 +185,8 @@ cdef class Curve(object):
                     cursor += n * sizeof(TRatePt)
                     src += n * (sizeof(double) + sizeof(uint16_t))
                     memcpy(cursor, src, decomp_size)
+            else:
+                raise ValueError("fmt types can be either Fmt.Plain, Fmt.LZ4 or Fmt.Packed")
             instance.buf_size = state_size
             instance.buf.reset(curve, char_free)
         return instance
@@ -779,8 +780,8 @@ cdef class SpreadCurve(Curve):
             memcpy(cursor, &self.defaulted, sizeof(TDate))
             cursor += sizeof(TDate)
             self.offset_name = cursor - self.buf.get()
-            cursor[0] =<char><CurveName.Seniority>seniority
-            cursor[1] = <char><CurveName.DocClause>doc_clause
+            cursor[0] = <char>seniority
+            cursor[1] = <char>doc_clause
             strncpy(&cursor[2], c_ticker, ticker_len + 1)
 
     cdef inline double* recovery_rates_ptr(self) nogil:
@@ -868,7 +869,7 @@ cdef class SpreadCurve(Curve):
             curve.fArray[i].fRate = JPMCDS_MAX_RATE
 
     @classmethod
-    def from_bytes(cls, object state, const int fmt=0):
+    def from_bytes(cls, object state, Fmt fmt=Fmt.Plain):
         cdef:
             SpreadCurve instance = super().from_bytes(state, fmt)
             char* cursor = instance.buf.get()