Ensure that annotations are strings

friendly_states/core.py

@@ -410,6 +410,13 @@ def set_state(self, previous_state, new_state):
 
 
 def extract_state_names(annotation):
+    if not isinstance(annotation, str):
+        raise ValueError(
+            "Found non-string annotation. Remember to add:\n\n"
+            "from __future__ import annotations\n\n"
+            "at the top of your file."
+        )
+
     try:
         tree = ast.parse(annotation)
     except SyntaxError:

tests/test_core.py

@@ -6,7 +6,7 @@
 
 import pytest
 
-from friendly_states.core import AttributeState, IncorrectInitialState, BaseState, MappingKeyState
+from friendly_states.core import AttributeState, IncorrectInitialState, BaseState, MappingKeyState, extract_state_names
 from friendly_states.exceptions import StateChangedElsewhere, IncorrectSummary, MultipleMachineAncestors, \
     InheritedFromState, CannotInferOutputState, DuplicateStateNames, DuplicateOutputStates, UnknownOutputState, \
     ReturnedInvalidState, GetStateDidNotReturnState
@@ -587,3 +587,13 @@ class Summary:
         CommonState2: []
 
     machine2.Machine.check_summary(Summary)
+
+
+def test_extract_state_names():
+    assert extract_state_names("x x") is None
+    assert extract_state_names("") is None
+    assert extract_state_names("x;x") is None
+    assert extract_state_names("[x[y]]") is None
+    assert extract_state_names("[x[y]]") is None
+    with raises(ValueError):
+        extract_state_names(None)