Making blocks more generic

src/omlt/block.py

@@ -32,18 +32,7 @@ class is used in combination with a formulation object to construct the
 from omlt.base import DEFAULT_MODELING_LANGUAGE, OmltVarFactory
 
 
-@declare_custom_block(name="OmltBlock")
-class OmltBlockData(BlockData):
-    def __init__(self, component):
-        super().__init__(component)
-        self.__formulation = None
-        self.__input_indexes = None
-        self.__output_indexes = None
-        self._format = DEFAULT_MODELING_LANGUAGE
-
-    def set_format(self, lang):
-        self._format = lang
-
+class OmltBlockCore:
     def _setup_inputs_outputs(self, *, input_indexes, output_indexes):
         """Setup inputs and outputs.
 
@@ -117,3 +106,16 @@ def build_formulation(self, formulation, lang=None):
 
         # tell the formulation object to construct the necessary models
         self.__formulation._build_formulation()
+
+
+@declare_custom_block(name="OmltBlock")
+class OmltBlockData(_BlockData, OmltBlockCore):
+    def __init__(self, component):
+        super().__init__(component)
+        self.__formulation = None
+        self.__input_indexes = None
+        self.__output_indexes = None
+        self._format = DEFAULT_MODELING_LANGUAGE
+
+    def set_format(self, lang):
+        self._format = lang

tests/base/test_block.py

@@ -46,9 +46,9 @@ def test_block():
     formulation = DummyFormulation()
     m.b.build_formulation(formulation, lang="pyomo")
 
-    assert m.b._OmltBlockData__formulation is formulation
-    assert list(m.b.inputs) == ["A", "C", "D"]
-    assert list(m.b.outputs) == [(0, 0), (0, 1), (1, 0), (1, 1)]
+    assert m.b._OmltBlockCore__formulation is formulation
+    assert [k for k in m.b.inputs] == ["A", "C", "D"]
+    assert [k for k in m.b.outputs] == [(0, 0), (0, 1), (1, 0), (1, 1)]
 
 
 def test_input_output_auto_creation():