rod 0.2.1.dev45__tar.gz → 0.2.1.dev57__tar.gz

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: rod
-Version: 0.2.1.dev45
+Version: 0.2.1.dev57
 Summary: The ultimate Python tool for RObot Descriptions processing.
 Home-page: https://github.com/ami-iit/rod
 Author: Diego Ferigo

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/__init__.py

@@ -115,9 +115,9 @@ def check_compatible_sdformat(specification_version: str) -> None:
 
     if not GazeboHelper.has_gazebo():
         return
-    else:
-        cmdline = GazeboHelper.get_gazebo_executable()
-        logging.info(f"Calling sdformat through '{cmdline} sdf'")
+
+    cmdline = GazeboHelper.get_gazebo_executable()
+    logging.info(f"Calling sdformat through '{cmdline} sdf'")
 
     output_sdf_version = packaging.version.Version(
         xmltodict.parse(

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/builder/primitive_builder.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import abc
 import dataclasses
-from typing import Optional, Union
+from typing import Optional
 
 import numpy as np
 import numpy.typing as npt
@@ -14,7 +16,7 @@ class PrimitiveBuilder(abc.ABC):
     name: str
     mass: float
 
-    element: Union[rod.Model, rod.Link, rod.Inertial, rod.Collision, rod.Visual] = (
+    element: rod.Model | rod.Link | rod.Inertial | rod.Collision | rod.Visual = (
         dataclasses.field(
             default=None, init=False, repr=False, hash=False, compare=False
         )
@@ -22,7 +24,7 @@ class PrimitiveBuilder(abc.ABC):
 
     def build(
         self,
-    ) -> Union[rod.Model, rod.Link, rod.Inertial, rod.Collision, rod.Visual]:
+    ) -> rod.Model | rod.Link | rod.Inertial | rod.Collision | rod.Visual:
         return self.element
 
     # ================
@@ -43,9 +45,9 @@ class PrimitiveBuilder(abc.ABC):
 
     def build_model(
         self,
-        name: Optional[str] = None,
-        pose: Optional[rod.Pose] = None,
-    ) -> "PrimitiveBuilder":
+        name: str | None = None,
+        pose: rod.Pose | None = None,
+    ) -> PrimitiveBuilder:
         self._check_element()
 
         self.element = self._model(name=name, pose=pose)
@@ -54,16 +56,16 @@ class PrimitiveBuilder(abc.ABC):
 
     def build_link(
         self,
-        name: Optional[str] = None,
-        pose: Optional[rod.Pose] = None,
-    ) -> "PrimitiveBuilder":
+        name: str | None = None,
+        pose: rod.Pose | None = None,
+    ) -> PrimitiveBuilder:
         self._check_element()
 
         self.element = self._link(name=name, pose=pose)
 
         return self
 
-    def build_inertial(self, pose: Optional[rod.Pose] = None) -> "PrimitiveBuilder":
+    def build_inertial(self, pose: rod.Pose | None = None) -> PrimitiveBuilder:
         self._check_element()
 
         self.element = self._inertial(pose=pose)
@@ -72,9 +74,9 @@ class PrimitiveBuilder(abc.ABC):
 
     def build_visual(
         self,
-        name: Optional[str] = None,
-        pose: Optional[rod.Pose] = None,
-    ) -> "PrimitiveBuilder":
+        name: str | None = None,
+        pose: rod.Pose | None = None,
+    ) -> PrimitiveBuilder:
         self._check_element()
 
         self.element = self._visual(name=name, pose=pose)
@@ -83,9 +85,9 @@ class PrimitiveBuilder(abc.ABC):
 
     def build_collision(
         self,
-        name: Optional[str] = None,
-        pose: Optional[rod.Pose] = None,
-    ) -> "PrimitiveBuilder":
+        name: str | None = None,
+        pose: rod.Pose | None = None,
+    ) -> PrimitiveBuilder:
         self._check_element()
 
         self.element = self._collision(name=name, pose=pose)
@@ -98,10 +100,10 @@ class PrimitiveBuilder(abc.ABC):
 
     def add_link(
         self,
-        name: Optional[str] = None,
-        pose: Optional[rod.Pose] = None,
-        link: Optional[rod.Link] = None,
-    ) -> "PrimitiveBuilder":
+        name: str | None = None,
+        pose: rod.Pose | None = None,
+        link: rod.Link | None = None,
+    ) -> PrimitiveBuilder:
         if not isinstance(self.element, rod.Model):
             raise ValueError(type(self.element))
 
@@ -116,9 +118,9 @@ class PrimitiveBuilder(abc.ABC):
 
     def add_inertial(
         self,
-        pose: Optional[rod.Pose] = None,
-        inertial: Optional[rod.Inertial] = None,
-    ) -> "PrimitiveBuilder":
+        pose: rod.Pose | None = None,
+        inertial: rod.Inertial | None = None,
+    ) -> PrimitiveBuilder:
         if not isinstance(self.element, (rod.Model, rod.Link)):
             raise ValueError(type(self.element))
 
@@ -144,11 +146,11 @@ class PrimitiveBuilder(abc.ABC):
 
     def add_visual(
         self,
-        name: Optional[str] = None,
+        name: str | None = None,
         use_inertial_pose: bool = True,
-        pose: Optional[rod.Pose] = None,
-        visual: Optional[rod.Visual] = None,
-    ) -> "PrimitiveBuilder":
+        pose: rod.Pose | None = None,
+        visual: rod.Visual | None = None,
+    ) -> PrimitiveBuilder:
         if not isinstance(self.element, (rod.Model, rod.Link)):
             raise ValueError(type(self.element))
 
@@ -180,11 +182,11 @@ class PrimitiveBuilder(abc.ABC):
 
     def add_collision(
         self,
-        name: Optional[str] = None,
+        name: str | None = None,
         use_inertial_pose: bool = True,
-        pose: Optional[rod.Pose] = None,
-        collision: Optional[rod.Collision] = None,
-    ) -> "PrimitiveBuilder":
+        pose: rod.Pose | None = None,
+        collision: rod.Collision | None = None,
+    ) -> PrimitiveBuilder:
         if not isinstance(self.element, (rod.Model, rod.Link)):
             raise ValueError(type(self.element))
 
@@ -224,8 +226,8 @@ class PrimitiveBuilder(abc.ABC):
 
     def _model(
         self,
-        name: Optional[str] = None,
-        pose: Optional[rod.Pose] = None,
+        name: str | None = None,
+        pose: rod.Pose | None = None,
     ) -> rod.Model:
         name = name if name is not None else self.name
         logging.debug(f"Building model '{name}'")
@@ -240,15 +242,15 @@ class PrimitiveBuilder(abc.ABC):
 
     def _link(
         self,
-        name: Optional[str] = None,
-        pose: Optional[rod.Pose] = None,
+        name: str | None = None,
+        pose: rod.Pose | None = None,
     ) -> rod.Link:
         return rod.Link(
             name=name if name is not None else f"{self.name}_link",
             pose=pose,
         )
 
-    def _inertial(self, pose: Optional[rod.Pose] = None) -> rod.Inertial:
+    def _inertial(self, pose: rod.Pose | None = None) -> rod.Inertial:
         return rod.Inertial(
             pose=pose,
             mass=self.mass,
@@ -257,8 +259,8 @@ class PrimitiveBuilder(abc.ABC):
 
     def _visual(
         self,
-        name: Optional[str] = None,
-        pose: Optional[rod.Pose] = None,
+        name: str | None = None,
+        pose: rod.Pose | None = None,
     ) -> rod.Visual:
         name = name if name is not None else f"{self.name}_visual"
 
@@ -271,7 +273,7 @@ class PrimitiveBuilder(abc.ABC):
     def _collision(
         self,
         name: Optional[str],
-        pose: Optional[rod.Pose] = None,
+        pose: rod.Pose | None = None,
     ) -> rod.Collision:
         name = name if name is not None else f"{self.name}_collision"
 
@@ -297,7 +299,7 @@ class PrimitiveBuilder(abc.ABC):
         relative_to: str = None,
         degrees: bool = None,
         rotation_format: str = None,
-    ) -> Optional[rod.Pose]:
+    ) -> rod.Pose | None:
         if pos is None and rpy is None:
             return rod.Pose.from_transform(transform=np.eye(4), relative_to=relative_to)
 

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/builder/primitives.py

@@ -1,6 +1,5 @@
 import dataclasses
 import pathlib
-from typing import Union
 
 import trimesh
 from numpy.typing import NDArray
@@ -63,7 +62,7 @@ class CylinderBuilder(PrimitiveBuilder):
 
 @dataclasses.dataclass
 class MeshBuilder(PrimitiveBuilder):
-    mesh_path: Union[str, pathlib.Path]
+    mesh_path: str | pathlib.Path
     scale: NDArray
 
     def __post_init__(self) -> None:

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/kinematics/kinematic_tree.py

@@ -1,7 +1,9 @@
+from __future__ import annotations
+
 import copy
 import dataclasses
 import functools
-from typing import Dict, List, Sequence, Tuple, Union
+from typing import Dict, List, Sequence, Tuple
 
 import numpy as np
 
@@ -12,7 +14,7 @@ from rod.tree import DirectedTree, DirectedTreeNode, TreeEdge, TreeFrame
 
 @dataclasses.dataclass(frozen=True)
 class KinematicTree(DirectedTree):
-    model: "rod.Model"
+    model: rod.Model
 
     joints: List[TreeEdge] = dataclasses.field(default_factory=list)
     frames: List[TreeFrame] = dataclasses.field(default_factory=list)
@@ -46,7 +48,7 @@ class KinematicTree(DirectedTree):
         return [joint.name() for joint in self.joints]
 
     @staticmethod
-    def build(model: "rod.Model", is_top_level: bool = True) -> "KinematicTree":
+    def build(model: rod.Model, is_top_level: bool = True) -> KinematicTree:
         logging.debug(msg=f"Building kinematic tree of model '{model.name}'")
 
         if model.model is not None:
@@ -199,7 +201,7 @@ class KinematicTree(DirectedTree):
             new_base_node, additional_frames = KinematicTree.remove_edge(
                 edge=world_to_base_edge, keep_parent=False
             )
-            assert any([f.name() == TreeFrame.WORLD for f in additional_frames])
+            assert any(f.name() == TreeFrame.WORLD for f in additional_frames)
 
             # Replace the former base node with the new base node
             nodes_links_dict[new_base_node.name()] = new_base_node

rod-0.2.1.dev57/src/rod/kinematics/tree_transforms.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import copy
+import dataclasses
+
+import numpy as np
+import numpy.typing as npt
+
+import rod
+from rod.kinematics.kinematic_tree import KinematicTree
+from rod.tree import TreeFrame
+
+
+@dataclasses.dataclass
+class TreeTransforms:
+    kinematic_tree: KinematicTree = dataclasses.dataclass(init=False)
+    _transform_cache: dict[str, npt.NDArray] = dataclasses.field(default_factory=dict)
+
+    @staticmethod
+    def build(
+        model: rod.Model,
+        is_top_level: bool = True,
+    ) -> TreeTransforms:
+        model = copy.deepcopy(model)
+
+        # Make sure that all elements have a pose attribute with explicit 'relative_to'.
+        model.resolve_frames(is_top_level=is_top_level, explicit_frames=True)
+
+        # Build the kinematic tree and return the TreeTransforms object.
+        return TreeTransforms(
+            kinematic_tree=KinematicTree.build(model=model, is_top_level=is_top_level)
+        )
+
+    def transform(self, name: str) -> npt.NDArray:
+        if name in self._transform_cache:
+            return self._transform_cache[name]
+
+        self._transform_cache[name] = self._compute_transform(name=name)
+        return self._transform_cache[name]
+
+    def _compute_transform(self, name: str) -> npt.NDArray:
+        match name:
+            case TreeFrame.WORLD:
+
+                return np.eye(4)
+
+            case name if name in {TreeFrame.MODEL, self.kinematic_tree.model.name}:
+
+                relative_to = self.kinematic_tree.model.pose.relative_to
+                assert relative_to in {None, ""}, (relative_to, name)
+                return self.kinematic_tree.model.pose.transform()
+
+            case name if name in self.kinematic_tree.joint_names():
+
+                edge = self.kinematic_tree.joints_dict[name]
+                assert edge.name() == name
+
+                # Get the pose of the frame in which the node's pose is expressed
+                assert edge._source.pose.relative_to not in {"", None}
+                x_H_E = edge._source.pose.transform()
+                W_H_x = self.transform(name=edge._source.pose.relative_to)
+
+                # Compute the world-to-node transform
+                # TODO: this assumes all joint positions to be 0
+                W_H_E = W_H_x @ x_H_E
+
+                return W_H_E
+
+            case name if name in self.kinematic_tree.link_names():
+
+                element = self.kinematic_tree.links_dict[name]
+
+                assert element.name() == name
+                assert element._source.pose.relative_to not in {"", None}
+
+                # Get the pose of the frame in which the link's pose is expressed.
+                x_H_L = element._source.pose.transform()
+                W_H_x = self.transform(name=element._source.pose.relative_to)
+
+                # Compute the world transform of the link.
+                W_H_L = W_H_x @ x_H_L
+                return W_H_L
+
+            case name if name in self.kinematic_tree.frame_names():
+
+                element = self.kinematic_tree.frames_dict[name]
+
+                assert element.name() == name
+                assert element._source.pose.relative_to not in {"", None}
+
+                # Get the pose of the frame in which the frame's pose is expressed.
+                x_H_F = element._source.pose.transform()
+                W_H_x = self.transform(name=element._source.pose.relative_to)
+
+                # Compute the world transform of the frame.
+                W_H_F = W_H_x @ x_H_F
+                return W_H_F
+
+            case _:
+                raise ValueError(name)
+
+    def relative_transform(self, relative_to: str, name: str) -> npt.NDArray:
+
+        world_H_name = self.transform(name=name)
+        world_H_relative_to = self.transform(name=relative_to)
+
+        return TreeTransforms.inverse(world_H_relative_to) @ world_H_name
+
+    @staticmethod
+    def inverse(transform: npt.NDArray) -> npt.NDArray:
+
+        R = transform[0:3, 0:3]
+        p = np.vstack(transform[0:3, 3])
+
+        return np.block(
+            [
+                [R.T, -R.T @ p],
+                [0, 0, 0, 1],
+            ]
+        )

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/logging.py

@@ -1,6 +1,5 @@
 import enum
 import logging
-from typing import Union
 
 import coloredlogs
 
@@ -20,7 +19,7 @@ def _logger() -> logging.Logger:
     return logging.getLogger(name=LOGGER_NAME)
 
 
-def set_logging_level(level: Union[int, LoggingLevel] = LoggingLevel.WARNING):
+def set_logging_level(level: int | LoggingLevel = LoggingLevel.WARNING):
     if isinstance(level, int):
         level = LoggingLevel(level)
 

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/pretty_printer.py

@@ -30,7 +30,7 @@ class DataclassPrettyPrinter(abc.ABC):
         ]
 
         return (
-            f"[\n"
+            "[\n"
             + ",\n".join(f"{spacing_level}{el!s}" for el in list_str)
             + f",\n{spacing_level_up}]"
         )
@@ -45,26 +45,27 @@ class DataclassPrettyPrinter(abc.ABC):
         for field in dataclasses.fields(obj):
             attr = getattr(obj, field.name)
 
-            if attr is None or attr == "":
-                continue
-
-            elif isinstance(attr, list):
-                list_str = DataclassPrettyPrinter.list_to_string(
-                    obj=attr, level=level + 1
-                )
-                serialization += [(field.name, list_str)]
-                continue
-
-            elif dataclasses.is_dataclass(attr):
-                dataclass_str = DataclassPrettyPrinter.dataclass_to_str(
-                    obj=attr, level=level + 1
-                )
-                serialization += [(field.name, dataclass_str)]
-                continue
-
-            else:
-                serialization += [(field.name, f"{attr!s}")]
-                continue
+            match attr:
+                case None | "":
+                    continue
+
+                case list():
+                    list_str = DataclassPrettyPrinter.list_to_string(
+                        obj=attr, level=level + 1
+                    )
+                    serialization += [(field.name, list_str)]
+                    continue
+
+                case _ if dataclasses.is_dataclass(attr):
+                    dataclass_str = DataclassPrettyPrinter.dataclass_to_str(
+                        obj=attr, level=level + 1
+                    )
+                    serialization += [(field.name, dataclass_str)]
+                    continue
+
+                case _:
+                    serialization += [(field.name, f"{attr!s}")]
+                    continue
 
         spacing = " " * 4
         spacing_level = spacing * level

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/sdf/element.py

@@ -1,5 +1,5 @@
 import dataclasses
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List
 
 import mashumaro.config
 import mashumaro.mixins.dict
@@ -37,7 +37,7 @@ class Element(mashumaro.mixins.dict.DataClassDictMixin, DataclassPrettyPrinter):
         false_vals = {"0", "False", "false"}
         assert data in true_vals.union(false_vals)
 
-        return True if data in true_vals else False
+        return data in true_vals
 
     @staticmethod
     def serialize_float(data: float) -> str:
@@ -50,7 +50,7 @@ class Element(mashumaro.mixins.dict.DataClassDictMixin, DataclassPrettyPrinter):
         return " ".join(np.array(data, dtype=str))
 
     @staticmethod
-    def deserialize_list(data: str, length: Optional[int] = None) -> List[float]:
+    def deserialize_list(data: str, length: int | None = None) -> List[float]:
         assert isinstance(data, str)
         array = np.atleast_1d(np.array(data.split(sep=" "), dtype=float).squeeze())
 

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/sdf/link.py

@@ -1,5 +1,5 @@
 import dataclasses
-from typing import List, Optional, Union
+from typing import List, Optional
 
 import mashumaro
 import numpy as np
@@ -73,11 +73,9 @@ class Link(Element):
 
     inertial: Optional[Inertial] = dataclasses.field(default=None)
 
-    visual: Optional[Union[Visual, List[Visual]]] = dataclasses.field(default=None)
+    visual: Optional[Visual | List[Visual]] = dataclasses.field(default=None)
 
-    collision: Optional[Union[Collision, List[Collision]]] = dataclasses.field(
-        default=None
-    )
+    collision: Optional[Collision | List[Collision]] = dataclasses.field(default=None)
 
     gravity: Optional[bool] = dataclasses.field(
         default=None,

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/sdf/model.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import dataclasses
-from typing import List, Optional, Union
+from typing import List, Optional
 
 import mashumaro
 
@@ -54,13 +56,13 @@ class Model(Element):
 
     pose: Optional[Pose] = dataclasses.field(default=None)
 
-    model: Optional[Union["Model", List["Model"]]] = dataclasses.field(default=None)
+    model: Optional[Model | List[Model]] = dataclasses.field(default=None)
 
-    frame: Optional[Union[Frame, List[Frame]]] = dataclasses.field(default=None)
+    frame: Optional[Frame | List[Frame]] = dataclasses.field(default=None)
 
-    link: Optional[Union[Link, List[Link]]] = dataclasses.field(default=None)
+    link: Optional[Link | List[Link]] = dataclasses.field(default=None)
 
-    joint: Optional[Union[Joint, List[Joint]]] = dataclasses.field(default=None)
+    joint: Optional[Joint | List[Joint]] = dataclasses.field(default=None)
 
     def is_fixed_base(self) -> bool:
         joints_having_world_parent = [j for j in self.joints() if j.parent == "world"]
@@ -80,7 +82,7 @@ class Model(Element):
 
         return self.links()[0].name
 
-    def models(self) -> List["Model"]:
+    def models(self) -> List[Model]:
         if self.model is None:
             return []
 
@@ -155,7 +157,7 @@ class Model(Element):
 
     def switch_frame_convention(
         self,
-        frame_convention: "rod.FrameConvention",
+        frame_convention: rod.FrameConvention,
         is_top_level: bool = True,
         explicit_frames: bool = True,
         attach_frames_to_links: bool = True,

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/sdf/sdf.py

@@ -1,7 +1,9 @@
+from __future__ import annotations
+
 import dataclasses
 import os
 import pathlib
-from typing import List, Optional, Union
+from typing import List, Optional
 
 import mashumaro
 import packaging.specifiers
@@ -19,9 +21,9 @@ from .world import World
 class Sdf(Element):
     version: str = dataclasses.field(metadata=mashumaro.field_options(alias="@version"))
 
-    world: Optional[Union[World, List[World]]] = dataclasses.field(default=None)
+    world: Optional[World | List[World]] = dataclasses.field(default=None)
 
-    model: Optional[Union[Model, List[Model]]] = dataclasses.field(default=None)
+    model: Optional[Model | List[Model]] = dataclasses.field(default=None)
 
     def worlds(self) -> List[World]:
         if self.world is None:
@@ -44,7 +46,7 @@ class Sdf(Element):
         return self.model
 
     @staticmethod
-    def load(sdf: Union[pathlib.Path, str], is_urdf: Optional[bool] = None) -> "Sdf":
+    def load(sdf: pathlib.Path | str, is_urdf: bool | None = None) -> Sdf:
         """
         Load an SDF resource.
 
@@ -80,7 +82,7 @@ class Sdf(Element):
             and len(sdf) <= MAX_PATH
             and pathlib.Path(sdf).is_file()
         ):
-            sdf_string = pathlib.Path(sdf).read_text()
+            sdf_string = pathlib.Path(sdf).read_text(encoding="utf-8")
             is_urdf = (
                 is_urdf if is_urdf is not None else pathlib.Path(sdf).suffix == ".urdf"
             )
@@ -99,8 +101,8 @@ class Sdf(Element):
         # Parse the SDF to dict
         try:
             xml_dict = xmltodict.parse(xml_input=sdf_string)
-        except Exception:
-            raise ValueError("Failed to parse 'sdf' argument")
+        except Exception as exc:
+            raise exc("Failed to parse 'sdf' argument")
 
         # Look for the top-level <sdf> element
         try:

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/sdf/world.py

@@ -1,5 +1,5 @@
 import dataclasses
-from typing import List, Optional, Union
+from typing import List, Optional
 
 import mashumaro
 
@@ -34,9 +34,9 @@ class World(Element):
 
     scene: Scene = dataclasses.field(default_factory=Scene)
 
-    model: Optional[Union[Model, List[Model]]] = dataclasses.field(default=None)
+    model: Optional[Model | List[Model]] = dataclasses.field(default=None)
 
-    frame: Optional[Union[Frame, List[Frame]]] = dataclasses.field(default=None)
+    frame: Optional[Frame | List[Frame]] = dataclasses.field(default=None)
 
     def models(self) -> List[Model]:
         if self.model is None:

{rod-0.2.1.dev45 → rod-0.2.1.dev57}/src/rod/tree/directed_tree.py

@@ -1,7 +1,7 @@
 import collections.abc
 import dataclasses
 import functools
-from typing import Any, Callable, Dict, Iterable, List, Optional, Union
+from typing import Any, Callable, Dict, Iterable, List
 
 from .tree_elements import DirectedTreeNode
 
@@ -33,7 +33,7 @@ class DirectedTree(collections.abc.Sequence):
     @staticmethod
     def breadth_first_search(
         root: DirectedTreeNode,
-        sort_children: Optional[Callable[[Any], Any]] = lambda node: node.name(),
+        sort_children: Callable[[Any], Any] | None = lambda node: node.name(),
     ) -> Iterable[DirectedTreeNode]:
         queue = [root]
 
@@ -69,8 +69,8 @@ class DirectedTree(collections.abc.Sequence):
         )
 
     def __getitem__(
-        self, key: Union[int, slice, str]
-    ) -> Union[DirectedTreeNode, List[DirectedTreeNode]]:
+        self, key: int | slice | str
+    ) -> DirectedTreeNode | List[DirectedTreeNode]:
         # Get the nodes' dictionary (already inserted in order following BFS)
         nodes_dict = self.nodes_dict
 
@@ -100,7 +100,7 @@ class DirectedTree(collections.abc.Sequence):
     def __reversed__(self) -> Iterable[DirectedTreeNode]:
         yield from reversed(self)
 
-    def __contains__(self, item: Union[str, DirectedTreeNode]) -> bool:
+    def __contains__(self, item: str | DirectedTreeNode) -> bool:
         if isinstance(item, str):
             return item in self.nodes_dict.keys()