Make csp enum base class opt out

Signed-off-by: Tim Paine <[email protected]>

ccflow/enums.py

@@ -15,6 +15,7 @@
 """
 
 import inspect
+from os import environ
 from typing import Any, Callable, Dict, List, Union
 
 from pydantic import GetJsonSchemaHandler
@@ -23,36 +24,35 @@
 
 __all__ = ("auto", "Enum", "make_enum")
 
-try:
-    from csp import DynamicEnum, Enum as BaseEnum
 
-    auto = BaseEnum.auto
+from enum import (
+    Enum as BaseEnum,  # noqa: F401
+    auto,
+)
 
-    _CSP_ENUM = True
+BaseEnum.auto = auto
 
-except ImportError:
+_CSP_ENUM = False
+
+if environ.get("CCFLOW_NO_CSP", "0") != "1":
     try:
-        # NOTE: we also except attribute errors
-        # because csp.trading depends on cubist-core,
-        # and during struct hint/generation csp's
-        # enum is not fully constructed yet
-        # CSP BaseEnum just uses auto
-        from enum import auto
+        from csp import DynamicEnum, Enum as BaseEnum
 
-        from csp.impl.enum import DynamicEnum, Enum as BaseEnum
+        auto = BaseEnum.auto
 
         _CSP_ENUM = True
 
     except ImportError:
-        # if csp is not installed, rely on python Enum
-        from enum import (
-            Enum as BaseEnum,  # noqa: F401
-            auto,
-        )
+        try:
+            from enum import auto
+
+            from csp.impl.enum import DynamicEnum, Enum as BaseEnum
 
-        BaseEnum.auto = auto
+            _CSP_ENUM = True
 
-        _CSP_ENUM = False
+        except ImportError:
+            # if csp is not installed, rely on python Enum
+            pass
 
 
 class Enum(BaseEnum):

ccflow/tests/enums/test_enums.py

@@ -1,5 +1,7 @@
 import enum
 import importlib
+import os
+import sys
 from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
@@ -14,6 +16,7 @@ def tearDown(self) -> None:
         # Because test_init_parent and test_init_parent_csp muck around with imports
         # Make sure we always rest the imports at the end of each test so that other
         # tests are unaffected
+        os.environ.pop("CCFLOW_NO_CSP", None)
         importlib.invalidate_caches()
         import ccflow.enums
 
@@ -44,7 +47,7 @@ def g():
             self.assertRaises(ImportError, f)
             self.assertRaises(ImportError, g)
 
-            import ccflow
+            import ccflow.enums
 
             importlib.reload(ccflow.enums)
 
@@ -106,3 +109,25 @@ class MyAutoEnum(Enum):
         self.assertEqual(MyDynamicEnum.A.value, 2)
         self.assertEqual(list(MyDynamicEnum), [MyDynamicEnum.A, MyDynamicEnum.B])
         self.assertTrue(issubclass(MyDynamicEnum, Enum))
+
+    def test_init_no_csp_explicit(self):
+        os.environ["CCFLOW_NO_CSP"] = "1"
+
+        sys.modules.pop("csp", None)
+        importlib.invalidate_caches()
+        import ccflow.enums
+
+        importlib.reload(ccflow.enums)
+
+        self.assertTrue("csp" not in sys.modules)
+
+        os.environ["CCFLOW_NO_CSP"] = "0"
+
+        sys.modules.pop("csp", None)
+        importlib.invalidate_caches()
+        import ccflow.enums
+
+        importlib.reload(ccflow.enums)
+
+        self.assertTrue("csp" in sys.modules)
+        os.environ.pop("CCFLOW_NO_CSP", None)