diff --git a/state/base.py b/state/base.py
index 06e792f..614ae9e 100644
--- a/state/base.py
+++ b/state/base.py
@@ -191,7 +191,7 @@ class State(ABC):
         pass
 
     @abstractmethod
-    def read_dict_keys(self, elem: Element) -> Optional[List[Any]]:
+    def read_dict_keys(self, elem: Element) -> Optional[List[Element]]:
         """
         Reads labels of outgoing edges starting in given node.
 
diff --git a/state/pystate.py b/state/pystate.py
index b88a358..0cc7f7a 100644
--- a/state/pystate.py
+++ b/state/pystate.py
@@ -112,7 +112,7 @@ class PyState(State):
         else:
             return self.edges[e][1]
 
-    def read_dict_keys(self, elem: Element) -> Optional[List[Any]]:
+    def read_dict_keys(self, elem: Element) -> Optional[List[Element]]:
         if elem not in self.nodes and elem not in self.edges:
             return None