Use overloads to make get_attr type safe

The overloads allow us to properly distinguish between:
- Data type passed or not passed (returns str)
- Optional set to true or false (returns Optional[T] or T)

Also remove the now redundant casts.

Author: haudren-woven

launch/launch/action.py

@@ -20,7 +20,6 @@
 from typing import Text
 from typing import Tuple
 from typing import Dict
-from typing import cast
 
 from .condition import Condition
 from .launch_context import LaunchContext
@@ -62,8 +61,8 @@ def parse(entity: 'Entity', parser: 'Parser'):
         # Import here for avoiding cyclic imports.
         from .conditions import IfCondition
         from .conditions import UnlessCondition
-        if_cond = cast(str, entity.get_attr('if', optional=True))
-        unless_cond = cast(str, entity.get_attr('unless', optional=True))
+        if_cond = entity.get_attr('if', optional=True)
+        unless_cond = entity.get_attr('unless', optional=True)
         kwargs: Dict[str, Condition] = {}
         if if_cond is not None and unless_cond is not None:
             raise RuntimeError("if and unless conditions can't be used simultaneously")

launch/launch/actions/append_environment_variable.py

@@ -17,7 +17,6 @@
 import os
 from typing import List
 from typing import Union
-from typing import cast
 
 from ..action import Action
 from ..frontend import Entity
@@ -76,14 +75,14 @@ def parse(
     ):
         """Parse an 'append_env' entity."""
         _, kwargs = super().parse(entity, parser)
-        kwargs['name'] = parser.parse_substitution(cast(str, entity.get_attr('name')))
-        kwargs['value'] = parser.parse_substitution(cast(str, entity.get_attr('value')))
+        kwargs['name'] = parser.parse_substitution(entity.get_attr('name'))
+        kwargs['value'] = parser.parse_substitution(entity.get_attr('value'))
         prepend = entity.get_attr('prepend', optional=True, data_type=bool, can_be_str=True)
         if prepend is not None:
-            kwargs['prepend'] = parser.parse_if_substitutions(cast(Union[str, bool], prepend))
+            kwargs['prepend'] = parser.parse_if_substitutions(prepend)
         separator = entity.get_attr('separator', optional=True)
         if separator is not None:
-            kwargs['separator'] = parser.parse_substitution(cast(str, separator))
+            kwargs['separator'] = parser.parse_substitution(separator)
         return cls, kwargs
 
     @property

launch/launch/actions/declare_launch_argument.py

@@ -17,7 +17,6 @@
 from typing import List
 from typing import Optional
 from typing import Text
-from typing import cast
 
 import launch.logging
 
@@ -166,19 +165,17 @@ def parse(
     ):
         """Parse `arg` tag."""
         _, kwargs = super().parse(entity, parser)
-        kwargs['name'] = parser.escape_characters(cast(str, entity.get_attr('name')))
+        kwargs['name'] = parser.escape_characters(entity.get_attr('name'))
         default_value = entity.get_attr('default', optional=True)
         if default_value is not None:
-            kwargs['default_value'] = parser.parse_substitution(cast(str, default_value))
+            kwargs['default_value'] = parser.parse_substitution(default_value)
         description = entity.get_attr('description', optional=True)
         if description is not None:
-            kwargs['description'] = parser.escape_characters(cast(str, description))
-        # TODO: What to do here? get_attr is supposed to parse scalar / list of scalars, but in
-        # this case asks for a list of entities?
-        choices = cast(List[Entity], entity.get_attr('choice', data_type=List[Entity], optional=True))  # type: ignore
+            kwargs['description'] = parser.escape_characters(description)
+        choices = entity.get_attr('choice', data_type=List[Entity], optional=True)
         if choices is not None:
             kwargs['choices'] = [
-                parser.escape_characters(cast(str, choice.get_attr('value'))) for choice in choices
+                parser.escape_characters(choice.get_attr('value')) for choice in choices
             ]
         return cls, kwargs
 

launch/launch/actions/execute_process.py

@@ -318,35 +318,35 @@ def parse(
             ignore = []
 
         if 'cmd' not in ignore:
-            kwargs['cmd'] = cls._parse_cmdline(cast(str, entity.get_attr('cmd')), parser)
+            kwargs['cmd'] = cls._parse_cmdline(entity.get_attr('cmd'), parser)
 
         if 'cwd' not in ignore:
-            cwd = cast(Optional[str], entity.get_attr('cwd', optional=True))
+            cwd = entity.get_attr('cwd', optional=True)
             if cwd is not None:
                 kwargs['cwd'] = parser.parse_substitution(cwd)
 
         if 'name' not in ignore:
-            name = cast(Optional[str], entity.get_attr('name', optional=True))
+            name = entity.get_attr('name', optional=True)
             if name is not None:
                 kwargs['name'] = parser.parse_substitution(name)
 
         if 'prefix' not in ignore:
-            prefix = cast(Optional[str], entity.get_attr('launch-prefix', optional=True))
+            prefix = entity.get_attr('launch-prefix', optional=True)
             if prefix is not None:
                 kwargs['prefix'] = parser.parse_substitution(prefix)
 
         if 'output' not in ignore:
-            output = cast(Optional[str], entity.get_attr('output', optional=True))
+            output = entity.get_attr('output', optional=True)
             if output is not None:
                 kwargs['output'] = parser.parse_substitution(output)
 
         if 'respawn' not in ignore:
-            respawn = cast(Optional[str], entity.get_attr('respawn', optional=True))
+            respawn = entity.get_attr('respawn', optional=True)
             if respawn is not None:
                 kwargs['respawn'] = parser.parse_substitution(respawn)
 
         if 'respawn_delay' not in ignore:
-            respawn_delay = cast(Optional[float], entity.get_attr('respawn_delay', data_type=float, optional=True))
+            respawn_delay = entity.get_attr('respawn_delay', data_type=float, optional=True)
             if respawn_delay is not None:
                 if respawn_delay < 0.0:
                     raise ValueError(
@@ -356,7 +356,7 @@ def parse(
                 kwargs['respawn_delay'] = respawn_delay
 
         if 'sigkill_timeout' not in ignore:
-            sigkill_timeout = cast(Optional[float], entity.get_attr('sigkill_timeout', data_type=float, optional=True))
+            sigkill_timeout = entity.get_attr('sigkill_timeout', data_type=float, optional=True)
             if sigkill_timeout is not None:
                 if sigkill_timeout < 0.0:
                     raise ValueError(
@@ -366,7 +366,7 @@ def parse(
                 kwargs['sigkill_timeout'] = str(sigkill_timeout)
 
         if 'sigterm_timeout' not in ignore:
-            sigterm_timeout = cast(Optional[float], entity.get_attr('sigterm_timeout', data_type=float, optional=True))
+            sigterm_timeout = entity.get_attr('sigterm_timeout', data_type=float, optional=True)
             if sigterm_timeout is not None:
                 if sigterm_timeout < 0.0:
                     raise ValueError(
@@ -376,25 +376,24 @@ def parse(
                 kwargs['sigterm_timeout'] = str(sigterm_timeout)
 
         if 'shell' not in ignore:
-            shell = cast(Optional[bool], entity.get_attr('shell', data_type=bool, optional=True))
+            shell = entity.get_attr('shell', data_type=bool, optional=True)
             if shell is not None:
                 kwargs['shell'] = shell
 
         if 'emulate_tty' not in ignore:
-            emulate_tty = cast(Optional[bool], entity.get_attr('emulate_tty', data_type=bool, optional=True))
+            emulate_tty = entity.get_attr('emulate_tty', data_type=bool, optional=True)
             if emulate_tty is not None:
                 kwargs['emulate_tty'] = emulate_tty
 
         if 'additional_env' not in ignore:
             # Conditions won't be allowed in the `env` tag.
             # If that feature is needed, `set_enviroment_variable` and
             # `unset_enviroment_variable` actions should be used.
-            # TODO: Fixup the data_type annotation
-            env = cast(Optional[List[Entity]], entity.get_attr('env', data_type=List[Entity], optional=True))  # type: ignore
+            env = entity.get_attr('env', data_type=List[Entity], optional=True)
             if env is not None:
                 kwargs['additional_env'] = {
-                    tuple(parser.parse_substitution(cast(str, e.get_attr('name')))):
-                    parser.parse_substitution(cast(str, e.get_attr('value'))) for e in env
+                    tuple(parser.parse_substitution(e.get_attr('name'))):
+                    parser.parse_substitution(e.get_attr('value')) for e in env
                 }
                 for e in env:
                     e.assert_entity_completely_parsed()

launch/launch/frontend/entity.py

@@ -15,14 +15,20 @@
 """Module for Entity class."""
 
 from typing import List
+from typing import Literal
 from typing import Optional
 from typing import Text
+from typing import Type
 from typing import Union
+from typing import TypeVar
+from typing import overload
 
 from launch.utilities.type_utils import AllowedTypesType
 from launch.utilities.type_utils import AllowedValueType
 
 
+TargetType = TypeVar("TargetType")
+
 class Entity:
     """Single item in the intermediate front_end representation."""
 
@@ -41,17 +47,46 @@ def children(self) -> List['Entity']:
         """Get the Entity's children."""
         raise NotImplementedError()
 
+    # We need a few overloads for type checking:
+    # - Depending on optional, the return value is either T or Optional[T].
+    # Unfortunately, if the optional is not present, we need another overload to denote the
+    # default, so this has three values: true, false and not present.
+    # - If no data type is passed, the return type is str. Similar to the above, it has two
+    # possibilities: present or not present.
+    # => 6 overloads to cover every combination
+    @overload
+    def get_attr(self, name: Text, *, data_type: Type[TargetType], optional: Literal[False], can_be_str: bool = True) -> TargetType:
+        ...
+
+    @overload
+    def get_attr(self, name: Text, *, data_type: Type[TargetType], optional: Literal[True], can_be_str: bool = True) -> Optional[TargetType]:
+        ...
+
+    @overload
+    def get_attr(self, name: Text, *, data_type: Type[TargetType], can_be_str: bool = True) -> TargetType:
+        ...
+
+    @overload
+    def get_attr(self, name: Text, *, optional: Literal[False], can_be_str: bool = True) -> str:
+        ...
+
+    @overload
+    def get_attr(self, name: Text, *, optional: Literal[True], can_be_str: bool = True) -> Optional[str]:
+        ...
+
+    @overload
+    def get_attr(self, name: Text, *, can_be_str: bool = True) -> str:
+        ...
+
+
     def get_attr(
         self,
-        name: Text,
+        name,
         *,
-        data_type: AllowedTypesType = str,
-        optional: bool = False,
-        can_be_str: bool = True,
-    ) -> Optional[Union[
-        AllowedValueType,
-        List['Entity'],
-    ]]:
+        data_type=str,
+        optional=False,
+        can_be_str=True,
+    ):
         """
         Access an attribute of the entity.