fix primitive support

Author: markhalka

numba_cfunc_compiler/defaults/primitive_support.py

@@ -81,7 +81,22 @@ def try_parse_state(cls, node: ast.AnnAssign, var_name: str, globalns: dict) ->
         if not isinstance(node.value, ast.Constant):
             raise TypeError(f"State '{var_name}' must have a literal initial value")
 
-        return StateVariableInfo(var_name, node.value.value, state_type)
+        initial_value = node.value.value
+        initial_type = type(initial_value)
+
+        # State storage is allocated by the host from the concrete Python value,
+        # while generated code reads it using the declared State type. Keep those
+        # representations identical, allowing only the safe numeric widening that
+        # is already supported for inputs.
+        if state_type is float and initial_type is int:
+            initial_value = float(initial_value)
+        elif initial_type is not state_type:
+            raise TypeError(
+                f"State '{var_name}' expected an initial value of type "
+                f"{state_type.__name__}, got {initial_type.__name__}"
+            )
+
+        return StateVariableInfo(var_name, initial_value, state_type)
 
 
 def register():

numba_cfunc_compiler/tests/test_support_units.py

@@ -475,6 +475,11 @@ def test_models_type_factory_registry_and_source_registry():
         assert TypeFactory.try_parse_input(param, int)[1] == ParameterInfo(int)
         assert TypeFactory.try_parse_input(param, str) is None
         assert TypeFactory.try_parse_state(parse_stmt("x: State[int] = 1"), "x", {}) == StateVariableInfo("x", 1, int)
+        assert TypeFactory.try_parse_state(parse_stmt("x: State[float] = 1"), "x", {}) == StateVariableInfo("x", 1.0, float)
+        assert TypeFactory.try_parse_state(parse_stmt("x: State[bool] = True"), "x", {}) == StateVariableInfo("x", True, bool)
+        for annotation, value in (("int", "1.0"), ("int", "True"), ("float", "True"), ("bool", "1")):
+            with pytest.raises(TypeError, match=f"expected an initial value of type {annotation}"):
+                TypeFactory.try_parse_state(parse_stmt(f"x: State[{annotation}] = {value}"), "x", {})
         assert TypeFactory.try_parse_state(parse_stmt("x: State[str] = 'a'"), "x", {}) is None
 
         int_info = NumbaTypeRegistry.get_by_python_type(int)