Phase one of csp.Enum retirement: add support for IntEnum as native csp types

Signed-off-by: Rob Ambalu <robert.ambalu@point72.com>

Author: robambalu

cpp/csp/python/Conversions.h

@@ -446,6 +446,16 @@ inline PyObject * toPython( const CspEnum & e, const CspType & type )
     auto & enumType = static_cast<const CspEnumType&>( type );
     const auto * emeta = static_cast<const DialectCspEnumMeta*>( enumType.meta().get() );
 
+    if( emeta -> isPyIntEnum() )
+    {
+        //TODO - precache IntEnums on DialectCspEnumMeta instead of PyCspEnumMeta for faster IntEnum conversion
+        PyObjectPtr val = PyObjectPtr::own( toPython( e.value() ) );
+        PyObject * obj = PyObject_CallOneArg( ( PyObject * ) emeta -> pyType().get(), val.get() );
+        if( !obj )
+            CSP_THROW( PythonPassthrough, "" );
+        return obj;
+    }
+    
     PyObject * obj = emeta -> pyMeta() -> toPyEnum( e );
     if( !obj ) [[unlikely]]
         CSP_THROW( ValueError, e.value() << " is not a valid value on csp.enum type " << emeta -> name() );
@@ -457,9 +467,14 @@ inline CspEnum fromPython( PyObject * o, const CspType & type )
 {
     assert( type.type() == CspType::Type::ENUM );
 
-    if( !PyType_IsSubtype( Py_TYPE( o ), &PyCspEnum::PyType ) ||
-        static_cast<PyCspEnum *>( o ) -> meta() != static_cast<const CspEnumType &>( type ).meta().get() )
-        CSP_THROW( TypeError, "Invalid enum type, expected enum type " << static_cast<const CspEnumType &>( type ).meta() -> name() << " got " << Py_TYPE( o ) -> tp_name );
+    auto & enumType = static_cast<const CspEnumType&>( type );
+    const auto * emeta = static_cast<const DialectCspEnumMeta*>( enumType.meta().get() );
+    
+    if( !PyObject_IsInstance( o, ( PyObject * ) emeta -> pyType().get() ) )
+        CSP_THROW( TypeError, "Invalid enum type, expected enum type " << emeta -> pyType() -> tp_name << " got " << Py_TYPE( o ) -> tp_name );
+    
+    if( emeta -> isPyIntEnum() )
+        return static_cast<const CspEnumType &>( type ).meta() -> create( PyLong_AsLong( o ) );
 
     return static_cast<PyCspEnum *>( o ) -> enum_;
 }

cpp/csp/python/CspTypeFactory.cpp

@@ -7,6 +7,12 @@
 namespace csp::python
 {
 
+CspTypeFactory::CspTypeFactory()
+{
+    PyObject *enum_mod = PyImport_ImportModule( "enum" );
+    m_intEnumPyType    = PyTypeObjectPtr::own( ( PyTypeObject * ) PyObject_GetAttrString( enum_mod, "IntEnum" ) );
+}
+
 CspTypeFactory & CspTypeFactory::instance()
 {
     static CspTypeFactory s_instance;
@@ -64,6 +70,11 @@ CspTypePtr & CspTypeFactory::typeFromPyType( PyObject * pyTypeObj )
             auto meta = ( ( PyCspEnumMeta * ) pyType ) -> enumMeta;
             rv.first -> second = std::make_shared<csp::CspEnumType>( meta );
         }
+        else if( PyType_IsSubtype( pyType, m_intEnumPyType.get() ) )
+        {
+            auto meta = createCspEnumMetaFromIntEnum( PyTypeObjectPtr::incref( pyType ) );
+            rv.first -> second = std::make_shared<csp::CspEnumType>( meta );
+        }
         else if( pyType == PyDateTimeAPI -> DateTimeType )
             rv.first -> second = csp::CspType::DATETIME();
         else if( pyType == PyDateTimeAPI -> DeltaType )
@@ -88,4 +99,30 @@ void CspTypeFactory::removeCachedType( PyTypeObject * pyType )
     m_cache.erase( pyType );
 }
 
+std::shared_ptr<CspEnumMeta> CspTypeFactory::createCspEnumMetaFromIntEnum( PyTypeObjectPtr pyIntEnumType )
+{
+    CspEnumMeta::ValueDef metadef;
+    
+    PyObjectPtr iter = PyObjectPtr::check( PyObject_GetIter( ( PyObject * ) pyIntEnumType.get() ) );
+    PyObject * member;
+    while( ( member = PyIter_Next( iter.get() ) ) != NULL )
+    {
+        PyObjectPtr name = PyObjectPtr::check( PyObject_GetAttrString( member, "name" ) );
+
+        const char * namestr = PyUnicode_AsUTF8( name.get() );
+        if( !namestr )
+            CSP_THROW( PythonPassthrough, "" );
+            
+        if( !PyLong_Check( member ) )
+            CSP_THROW( TypeError, "enum key " << namestr << " expected an integer got " << PyObjectPtr::incref( member ) );
+        
+        int64_t value = fromPython<int64_t>( member );
+        metadef[ namestr ] = value;
+        
+        Py_DECREF( member );
+    }
+    
+    return std::make_shared<DialectCspEnumMeta>( pyIntEnumType, pyIntEnumType -> tp_name, metadef );
+}
+
 }

cpp/csp/python/CspTypeFactory.h

@@ -3,6 +3,7 @@
 
 #include <csp/core/Platform.h>
 #include <csp/engine/CspType.h>
+#include <csp/python/PyObjectPtr.h>
 #include <unordered_map>
 #include <Python.h>
 
@@ -17,9 +18,17 @@ class CSPTYPESIMPL_EXPORT CspTypeFactory
     CspTypePtr & typeFromPyType( PyObject * );
     void removeCachedType( PyTypeObject * );
 
+    PyTypeObject * intEnumPyType() { return m_intEnumPyType.get(); }
+
 private:
     using Cache = std::unordered_map<PyTypeObject *, CspTypePtr>;
+
+    std::shared_ptr<CspEnumMeta> createCspEnumMetaFromIntEnum( PyTypeObjectPtr pyIntEnumType );
+
+    CspTypeFactory();
     Cache m_cache;
+
+    PyTypeObjectPtr m_intEnumPyType;
 };
 
 }

cpp/csp/python/PyCspEnum.cpp

@@ -12,6 +12,7 @@ DialectCspEnumMeta::DialectCspEnumMeta( PyTypeObjectPtr pyType, const std::strin
     CspEnumMeta( name, def ),
     m_pyType( pyType )
 {
+    m_isPyIntEnum = PyType_IsSubtype( pyType.get(), CspTypeFactory::instance().intEnumPyType() );
 }
 
 /*

cpp/csp/python/PyCspEnum.h

@@ -40,11 +40,13 @@ class CSPTYPESIMPL_EXPORT DialectCspEnumMeta : public CspEnumMeta
 
     const PyTypeObjectPtr & pyType() const { return m_pyType; }
 
-    const PyCspEnumMeta * pyMeta() const   { return ( const PyCspEnumMeta * ) m_pyType.get(); }
+    const PyCspEnumMeta * pyMeta() const   { assert( !m_isPyIntEnum ); return ( const PyCspEnumMeta * ) m_pyType.get(); }
 
+    bool isPyIntEnum() const               { return m_isPyIntEnum; }
 private:
 
     PyTypeObjectPtr m_pyType;
+    bool            m_isPyIntEnum;
 };
 
 struct CSPTYPESIMPL_EXPORT PyCspEnum : public PyObject

cpp/csp/python/PyInputProxy.cpp

@@ -239,8 +239,8 @@ PyObject *PyInputProxy::valuesAt( ValueType valueType, PyObject *startIndexArg,
                                   PyObject *endIndexPolicyArg ) const
 {
     int32_t startIndex, endIndex;
-    auto startPolicy = static_cast<autogen::TimeIndexPolicy>( static_cast<PyCspEnum *>( startIndexPolicyArg ) -> enum_ );
-    auto endPolicy = static_cast<autogen::TimeIndexPolicy>( static_cast<PyCspEnum *>( endIndexPolicyArg ) -> enum_ );
+    auto startPolicy = csp::autogen::TimeIndexPolicy::create( startIndexPolicyArg );
+    auto endPolicy   = csp::autogen::TimeIndexPolicy::create( endIndexPolicyArg );
 
     if( startIndexArg == Py_None )
         startIndex = 1 - ts() -> numTicks();
@@ -513,9 +513,9 @@ static inline PyObject * PyInputProxy_values_at_impl( ValueType valueType, PyInp
     PyObject * endIndexArg;
     PyObject * startExclusiveArg;
     PyObject * endExclusiveArg;
-    if( !PyArg_ParseTuple( args, "OOO!O!", &startIndexArg, &endIndexArg,
-                           &PyCspEnum::PyType, &startExclusiveArg,
-                           &PyCspEnum::PyType, &endExclusiveArg ) )
+    if( !PyArg_ParseTuple( args, "OOOO", &startIndexArg, &endIndexArg,
+                           &startExclusiveArg,
+                           &endExclusiveArg ) )
         CSP_THROW( RuntimeException, "Invalid arguments passed to values_at" );
 
     return proxy -> valuesAt( valueType, startIndexArg, endIndexArg, startExclusiveArg, endExclusiveArg );

csp/build/csp_autogen.py

@@ -5,6 +5,7 @@
 import os.path
 import sys
 import types
+from enum import IntEnum
 
 # We need to patch mock modules into sys.modules to avoid pulling in csp/__init__.py and all of its baggage, which include imports of
 # _cspimpl, which would be a circular dep
@@ -85,7 +86,7 @@ def __init__(self, module_name: str, output_filename: str, namespace: str, gener
 
             if issubclass(v, Struct) and v is not Struct:
                 self._struct_types.append(v)
-            elif issubclass(v, Enum) and v is not Enum:
+            elif issubclass(v, Enum) and v is not Enum or issubclass(v, IntEnum) and v is not IntEnum:
                 self._enum_types.append(v)
 
     def _get_dependent_headers(self):
@@ -119,11 +120,9 @@ def cpp_filename(self):
         return self._cpp_filename
 
     def generate_header_code(self):
-        include_guard = "_IN_CSP_AUTOGEN_" + self._module_name.replace(".", "_").upper()
-        out = f"""
-#ifndef {include_guard}
-#define {include_guard}
-
+        out = f"""//// Generated from {self._module_name}
+#pragma once
+        
 """
         out += self._generate_headers()
 
@@ -139,11 +138,11 @@ def generate_header_code(self):
         for struct_type in self._struct_types:
             out += self._generate_struct_class(struct_type)
 
-        out += "\n}\n#endif"
+        out += "\n}"
         return out
 
     def _generate_headers(self):
-        common_headers = ["csp/core/Exception.h", "csp/engine/Struct.h", "cstddef"]
+        common_headers = ["csp/core/Exception.h", "csp/engine/Struct.h", "csp/python/Conversions.h", "cstddef"]
 
         common_headers.extend(self._get_dependent_headers())
         return "\n".join(f"#include <{h}>" for h in common_headers)
@@ -173,7 +172,11 @@ class {enum_name} : public csp::CspEnum
     static {enum_name} create( enum_ v )          {{ return s_meta -> create( ( int64_t ) v ); }}
     static {enum_name} create( const char * name) {{ return s_meta -> fromString( name ); }}
     static {enum_name} create( const std::string & s ) {{ return create( s.c_str() ); }}
-
+    static {enum_name} create( PyObject * e )
+    {{
+        return {enum_name}( csp::python::fromPython<CspEnum>( e, *s_cspEnumType ) );
+    }}
+    
     enum_ enum_value() const {{ return ( enum_ ) value(); }}
 
     static constexpr uint32_t num_types() {{ return {len([x for x in enum_type])}; }}
@@ -183,7 +186,7 @@ class {enum_name} : public csp::CspEnum
     {enum_name}( const csp::CspEnum & v ) : csp::CspEnum( v ) {{ CSP_TRUE_OR_THROW( v.meta() == s_meta.get(), AssertionError, "Mismatched enum meta" ); }}
 
 private:
-
+    static std::shared_ptr<const csp::CspEnumType> s_cspEnumType;
     static std::shared_ptr<csp::CspEnumMeta> s_meta;
 }};
 """
@@ -439,14 +442,18 @@ def generate_cpp_code(self):
         assert_or_die( enumType != nullptr, "failed to find num type {enum_name} in module {self._module_name}" );
 
         // should add some assertion here..
-        csp::python::PyCspEnumMeta * pymeta = ( csp::python::PyCspEnumMeta * ) enumType;
-        s_meta = pymeta -> enumMeta;
+        //csp::python::PyCspEnumMeta * pymeta = ( csp::python::PyCspEnumMeta * ) enumType;
+        //s_meta = pymeta -> enumMeta;
+        auto type = csp::python::CspTypeFactory::instance().typeFromPyType( enumType );
+        s_cspEnumType = std::static_pointer_cast<const CspEnumType>( type );
+        s_meta = s_cspEnumType -> meta();
     }}
 
     return true;
 }}
 
 bool static_init_{enum_name} = {enum_name}::static_init();
+std::shared_ptr<const csp::CspEnumType> {enum_name}::s_cspEnumType;
 std::shared_ptr<csp::CspEnumMeta> {enum_name}::s_meta;
 {static_decls}
 """
@@ -460,6 +467,7 @@ def generate_cpp_code(self):
 #include <csp/python/Common.h>
 #include <csp/python/PyStruct.h>
 #include <csp/python/PyCspEnum.h>
+#include <csp/python/CspTypeFactory.h>
 #include <iostream>
 #include <stdlib.h>
 #include <Python.h>

csp/impl/types/autogen_types.py

@@ -1,5 +1,6 @@
 # Set of structs / enums that are autogenerated / used in the C++ engine
-from csp.impl.enum import Enum
+from enum import IntEnum
+
 from csp.impl.struct import Struct
 
 CSP_AUTOGEN_HINTS = {"cpp_header": "csp/engine/csp_autogen/autogen_types.h"}
@@ -23,7 +24,7 @@ def removed(self):
         return [event.key for event in self.events if not event.added]
 
 
-class TimeIndexPolicy(Enum):
+class TimeIndexPolicy(IntEnum):
     """An enum that specifies the policy for handling the start and end values in functions like values_at."""
 
     INCLUSIVE = 1

csp/tests/test_history.py

@@ -161,7 +161,7 @@ def collect_index(
                 with self.assertRaises(TypeError):
                     csp.values_at(x, 0, 0, csp.TimeIndexPolicy.EXCLUSIVE, csp.TimeIndexPolicy.EXCLUSIVE)
 
-                with self.assertRaises(RuntimeError):
+                with self.assertRaises(TypeError):
                     csp.items_at(x, -10, 0, False, False)
 
         @csp.node
@@ -200,7 +200,7 @@ def collect_timedelta(
 
                 csp.output(values_default, csp.values_at(x))
 
-                with self.assertRaises(RuntimeError):
+                with self.assertRaises(TypeError):
                     csp.times_at(x, startIndex, endIndex, "abc", False)
 
         @csp.node