Fix(typing)

aiplan4eu · Apr 11, 2024 · d686e33 · d686e33
1 parent b3dac95
commit d686e33
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/unified_planning/model/walkers/free_vars.py b/unified_planning/model/walkers/free_vars.py
@@ -17,7 +17,7 @@
 import unified_planning.model.walkers as walkers
 from unified_planning.model.fnode import FNode
 from unified_planning.model.operators import OperatorKind
-from typing import List, FrozenSet
+from typing import Iterator, List, FrozenSet
 
 
 class FreeVarsExtractor(walkers.dag.DagWalker):
@@ -39,7 +39,7 @@ def get(self, expression: FNode) -> FrozenSet[FNode]:
     def walk_all_types(
         self, expression: FNode, args: List[FrozenSet[FNode]]
     ) -> FrozenSet[FNode]:
-        res_generator = (x for y in args for x in y)
+        res_generator: Iterator[FNode] = (x for y in args for x in y)
         if expression.is_fluent_exp():
             res_generator = chain(res_generator, (expression,))
         return frozenset(res_generator)
diff --git a/unified_planning/model/walkers/simplifier.py b/unified_planning/model/walkers/simplifier.py
@@ -15,7 +15,7 @@
 
 from fractions import Fraction
 from collections import OrderedDict
-from typing import Dict, List, Optional, Set, Union, cast
+from typing import List, Optional, FrozenSet, Union, cast
 import unified_planning as up
 import unified_planning.environment
 import unified_planning.model.walkers as walkers
@@ -178,7 +178,7 @@ def walk_implies(self, expression: FNode, args: List[FNode]) -> FNode:
 
     def walk_exists(self, expression: FNode, args: List[FNode]) -> FNode:
         assert len(args) == 1
-        free_vars: Set[
+        free_vars: FrozenSet[
             "up.model.variable.Variable"
         ] = self.environment.free_vars_oracle.get_free_variables(args[0])
         vars = set(var for var in expression.variables() if var in free_vars)
@@ -224,7 +224,7 @@ def walk_exists(self, expression: FNode, args: List[FNode]) -> FNode:
 
     def walk_forall(self, expression: FNode, args: List[FNode]) -> FNode:
         assert len(args) == 1
-        free_vars: Set[
+        free_vars: FrozenSet[
             "up.model.variable.Variable"
         ] = self.environment.free_vars_oracle.get_free_variables(args[0])
         vars = tuple(var for var in expression.variables() if var in free_vars)