Modernize type annotations in backend_reader.py and backend_writer.py.

envlogger/backends/backend_reader.py

@@ -17,7 +17,7 @@
 
 import abc
 from collections.abc import Callable, Iterator, Sequence
-from typing import Any, Generic, Optional, TypeVar, Union, overload
+from typing import Any, Generic, Self, TypeVar, overload
 
 from absl import logging
 from envlogger import step_data
@@ -50,7 +50,7 @@ def __getitem__(self, index: int) -> T:
   def __getitem__(self, index: slice) -> list[T]:
     ...
 
-  def __getitem__(self, index: Union[int, slice]) -> Union[T, list[T]]:
+  def __getitem__(self, index: int | slice) -> T | list[T]:
     """Retrieves items from this sequence.
 
     Args:
@@ -98,15 +98,15 @@ class BackendReader(metaclass=abc.ABCMeta):
   def __init__(self):
     self._init_visitors()
 
-  def copy(self) -> 'BackendReader':
+  def copy(self) -> Self:
     """Returns a copy of self."""
 
     c = self._copy()
     c._init_visitors()
     return c
 
   @abc.abstractmethod
-  def _copy(self) -> 'BackendReader':
+  def _copy(self) -> Self:
     """Implementation-specific copy behavior."""
 
   def _init_visitors(self):
@@ -151,12 +151,12 @@ def get_nth_step_from_episode(j: int):
     return _SequenceAdapter(
         count=episode.num_steps, get_nth_item=get_nth_step_from_episode)
 
-  def _get_nth_episode_metadata(self, i: int) -> Optional[Any]:
+  def _get_nth_episode_metadata(self, i: int) -> Any | None:
     """Returns the metadata for episode `i` (0-based)."""
     episode = self._get_nth_episode_info(i, include_metadata=True)
     return codec.decode(episode.metadata)
 
-  def __enter__(self):
+  def __enter__(self) -> Self:
     return self
 
   def __exit__(self, exc_type, exc_value, tb):
@@ -177,7 +177,7 @@ def metadata(self) -> dict[str, Any]:
   def episodes(self) -> Sequence[Sequence[step_data.StepData]]:
     return self._episodes
 
-  def episode_metadata(self) -> Sequence[Optional[Any]]:
+  def episode_metadata(self):
     return self._episode_metadata
 
   @property

envlogger/backends/backend_writer.py

@@ -16,7 +16,7 @@
 """Abstract trajectory logging interface."""
 
 import abc
-from typing import Any, Optional
+from typing import Any
 
 from envlogger import step_data
 from envlogger.backends import schedulers
@@ -27,8 +27,8 @@ class BackendWriter(metaclass=abc.ABCMeta):
 
   def __init__(
       self,
-      metadata: Optional[dict[str, Any]] = None,
-      scheduler: Optional[schedulers.Scheduler] = None,
+      metadata: dict[str, Any] | None = None,
+      scheduler: schedulers.Scheduler | None = None,
   ):
     """BackendWriter base class.