diff --git a/unified_planning/model/walkers/free_vars.py b/unified_planning/model/walkers/free_vars.py index 250efe1e8..a34faf857 100644 --- 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 index 83b0b4fcf..29ea78358 100644 --- 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)