# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/rlds_utils.py
# def to_rlds_step(prev_step: step_data.StepData,
#                  step: Optional[step_data.StepData]) -> Step:
# def _find_extra_shard(split_info: tfds.core.SplitInfo) -> Optional[Any]:
# def maybe_recover_last_shard(builder: tfds.core.DatasetBuilder):
# 
# Path: envlogger/backends/tfds_backend_testlib.py
# def generate_episode_data(
#     backend: backend_writer.BackendWriter,
#     num_episodes: int = 2,
# ) -> List[List[step_data.StepData]]:
# def catch_env_tfds_config(
#     name: str = 'catch_example') -> tfds.rlds.rlds_base.DatasetConfig:
# def tfds_backend_catch_env(
#     data_directory: str,
#     max_episodes_per_file: int = 1,
#     split_name: Optional[str] = None) -> tfds_backend_writer.TFDSBackendWriter:
# 
# Path: envlogger/backends/tfds_backend_writer.py
# class RLDSEnvloggerBuilder(tfds.core.GeneratorBasedBuilder):
# class Shard(object):
# class Episode(object):
# class Split(object):
# class TFDSBackendWriter(backend_writer.BackendWriter):
#   def _info(self) -> tfds.core.DatasetInfo:
#   def _split_generators(self, dl_manager: tfds.download.DownloadManager):
#   def _generate_examples(self, path):
#   def add_episode(self, serialized_episode: str) -> None:
#   def close_writer(self) -> None:
#   def add_step(self, step: step_data.StepData) -> None:
#   def serialize_episode(
#       self, features: tfds.features.FeaturesDict,
#       serializer: tfds.core.example_serializer.ExampleSerializer) -> str:
#   def update(self, shard: Shard) -> None:
#   def get_shard_path(self) -> str:
#   def get_split_dict(self) -> tfds.core.splits.SplitDict:
# def initialize_split(split_name: Optional[str], data_directory: Any,
#                      builder_name: str) -> Split:
#   def __init__(self,
#                data_directory: str,
#                ds_config: tfds.rlds.rlds_base.DatasetConfig,
#                max_episodes_per_file: int = 1000,
#                split_name: Optional[str] = None,
#                version: str = '0.0.1',
#                **base_kwargs):
#   def _write_split_metadata(self) -> None:
#   def _finalize_shard_and_update_split(self) -> None:
#   def _finalize_episode_and_update_shard(self) -> None:
#   def _write_and_reset_episode(self, is_last_episode: bool = False) -> None:
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Dict[str, Any]) -> None:
#   def close(self) -> None:


import os
import dm_env
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from typing import List
from absl.testing import absltest
from envlogger import step_data
from envlogger.backends import rlds_utils
from envlogger.backends import tfds_backend_testlib
from envlogger.backends import tfds_backend_writer
    for idx, rlds_step in enumerate(steps):
      step = expected_steps[idx + 1] if idx < len(expected_steps) - 1 else None
      expected_step = rlds_utils.to_rlds_step(expected_steps[idx], step)
      np.testing.assert_equal(expected_step, rlds_step)

  def test_backend_writer(self):
    num_episodes = 5
    max_episodes_per_file = 3
    data_dir = self.create_tempdir(name='my_data_dir').full_path
    expected_episodes = tfds_backend_testlib.generate_episode_data(
        backend=tfds_backend_testlib.tfds_backend_catch_env(
            data_directory=data_dir,
            max_episodes_per_file=max_episodes_per_file),
        num_episodes=num_episodes)

    builder = tfds.builder_from_directory(data_dir)
    ds = builder.as_dataset(split='train')

    num_episodes = 0
    for index, episode in enumerate(ds):
      self._assert_steps(expected_episodes[index], episode['steps'])
      self.assertEqual(episode['episode_id'], index)
      num_episodes += 1

    self.assertLen(expected_episodes, num_episodes)

  def test_backend_writer_with_split_name(self):
    num_episodes = 1
    max_episodes_per_file = 1
    data_dir = self.create_tempdir(name='my_data_dir').full_path

# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/rlds_utils.py
# def to_rlds_step(prev_step: step_data.StepData,
#                  step: Optional[step_data.StepData]) -> Step:
# def _find_extra_shard(split_info: tfds.core.SplitInfo) -> Optional[Any]:
# def maybe_recover_last_shard(builder: tfds.core.DatasetBuilder):
# 
# Path: envlogger/backends/tfds_backend_testlib.py
# def generate_episode_data(
#     backend: backend_writer.BackendWriter,
#     num_episodes: int = 2,
# ) -> List[List[step_data.StepData]]:
# def catch_env_tfds_config(
#     name: str = 'catch_example') -> tfds.rlds.rlds_base.DatasetConfig:
# def tfds_backend_catch_env(
#     data_directory: str,
#     max_episodes_per_file: int = 1,
#     split_name: Optional[str] = None) -> tfds_backend_writer.TFDSBackendWriter:
# 
# Path: envlogger/backends/tfds_backend_writer.py
# class RLDSEnvloggerBuilder(tfds.core.GeneratorBasedBuilder):
# class Shard(object):
# class Episode(object):
# class Split(object):
# class TFDSBackendWriter(backend_writer.BackendWriter):
#   def _info(self) -> tfds.core.DatasetInfo:
#   def _split_generators(self, dl_manager: tfds.download.DownloadManager):
#   def _generate_examples(self, path):
#   def add_episode(self, serialized_episode: str) -> None:
#   def close_writer(self) -> None:
#   def add_step(self, step: step_data.StepData) -> None:
#   def serialize_episode(
#       self, features: tfds.features.FeaturesDict,
#       serializer: tfds.core.example_serializer.ExampleSerializer) -> str:
#   def update(self, shard: Shard) -> None:
#   def get_shard_path(self) -> str:
#   def get_split_dict(self) -> tfds.core.splits.SplitDict:
# def initialize_split(split_name: Optional[str], data_directory: Any,
#                      builder_name: str) -> Split:
#   def __init__(self,
#                data_directory: str,
#                ds_config: tfds.rlds.rlds_base.DatasetConfig,
#                max_episodes_per_file: int = 1000,
#                split_name: Optional[str] = None,
#                version: str = '0.0.1',
#                **base_kwargs):
#   def _write_split_metadata(self) -> None:
#   def _finalize_shard_and_update_split(self) -> None:
#   def _finalize_episode_and_update_shard(self) -> None:
#   def _write_and_reset_episode(self, is_last_episode: bool = False) -> None:
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Dict[str, Any]) -> None:
#   def close(self) -> None:


# Path: envlogger/backends/tfds_backend_writer_test.py
import os
import dm_env
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from typing import List
from absl.testing import absltest
from envlogger import step_data
from envlogger.backends import rlds_utils
from envlogger.backends import tfds_backend_testlib
from envlogger.backends import tfds_backend_writer

    for idx, rlds_step in enumerate(steps):
      step = expected_steps[idx + 1] if idx < len(expected_steps) - 1 else None
      expected_step = rlds_utils.to_rlds_step(expected_steps[idx], step)
      np.testing.assert_equal(expected_step, rlds_step)

  def test_backend_writer(self):
    num_episodes = 5
    max_episodes_per_file = 3
    data_dir = self.create_tempdir(name='my_data_dir').full_path
    expected_episodes = tfds_backend_testlib.generate_episode_data(
        backend=tfds_backend_testlib.tfds_backend_catch_env(
            data_directory=data_dir,
            max_episodes_per_file=max_episodes_per_file),
        num_episodes=num_episodes)

    builder = tfds.builder_from_directory(data_dir)
    ds = builder.as_dataset(split='train')

    num_episodes = 0
    for index, episode in enumerate(ds):
      self._assert_steps(expected_episodes[index], episode['steps'])
      self.assertEqual(episode['episode_id'], index)
      num_episodes += 1

    self.assertLen(expected_episodes, num_episodes)

  def test_backend_writer_with_split_name(self):
    num_episodes = 1
    max_episodes_per_file = 1
    data_dir = self.create_tempdir(name='my_data_dir').full_path

    expected_episodes = tfds_backend_testlib.generate_episode_data(