# 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:


from absl.testing import absltest
from envlogger import step_data
from envlogger.backends import rlds_utils
from envlogger.backends import tfds_backend_testlib
import dm_env
import tensorflow_datasets as tfds

  def test_regenerate_splits_noop(self):
    num_episodes = 3
    max_episodes_per_file = 2
    data_dir = self.create_tempdir(name='my_data_dir').full_path
    _ = 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,
            split_name='split'),
        num_episodes=num_episodes)

    builder = tfds.builder_from_directory(data_dir)

    self.assertEqual(list(builder.info.splits.keys()), ['split'])
    self.assertEqual(builder.info.splits['split'].num_examples, 3)
    self.assertEqual(builder.info.splits['split'].num_shards, 2)
    self.assertEqual(builder.info.splits['split'].shard_lengths, [2, 1])

    expected_splits = builder.info.splits

    new_builder = rlds_utils.maybe_recover_last_shard(builder)

    self.assertEqual(new_builder.info.splits, expected_splits)

  def test_regenerate_ds_with_one_split(self):
    num_episodes = 3
    max_episodes_per_file = 5
    data_dir = self.create_tempdir(name='my_data_dir').full_path
    _ = tfds_backend_testlib.generate_episode_data(

# 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/rlds_utils_test.py
from absl.testing import absltest
from envlogger import step_data
from envlogger.backends import rlds_utils
from envlogger.backends import tfds_backend_testlib
import dm_env
import tensorflow_datasets as tfds


  def test_regenerate_splits_noop(self):
    num_episodes = 3
    max_episodes_per_file = 2
    data_dir = self.create_tempdir(name='my_data_dir').full_path
    _ = 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,
            split_name='split'),
        num_episodes=num_episodes)

    builder = tfds.builder_from_directory(data_dir)

    self.assertEqual(list(builder.info.splits.keys()), ['split'])
    self.assertEqual(builder.info.splits['split'].num_examples, 3)
    self.assertEqual(builder.info.splits['split'].num_shards, 2)
    self.assertEqual(builder.info.splits['split'].shard_lengths, [2, 1])

    expected_splits = builder.info.splits

    new_builder = rlds_utils.maybe_recover_last_shard(builder)

    self.assertEqual(new_builder.info.splits, expected_splits)

  def test_regenerate_ds_with_one_split(self):
    num_episodes = 3
    max_episodes_per_file = 5
    data_dir = self.create_tempdir(name='my_data_dir').full_path
    _ = tfds_backend_testlib.generate_episode_data(

        backend=tfds_backend_testlib.tfds_backend_catch_env(