# Path: envlogger/converters/codec.py
# _SCALAR_DIM_SIZE = -438
# def _python_int_to_bytes(py_int: int) -> bytes:
# def _set_datum_values_from_scalar(scalar: Union[ScalarNumber, bool, str, bytes],
#                                   datum: storage_pb2.Datum) -> bool:
# def _set_datum_values_from_array(array: np.ndarray,
#                                  values: storage_pb2.Datum.Values) -> None:
# def encode(
#     user_data: Union[np.ndarray, List[Any], Tuple[Any, ...], Dict[str, Any]]
# ) -> storage_pb2.Data:
# def decode_datum(
#     datum: storage_pb2.Datum
# ) -> Union[np.ndarray, ScalarNumber, bool, str, bytes]:
#     def from_bigint(int_bytes):
# def decode(
#     user_data: storage_pb2.Data
# ) -> Optional[Union[ScalarNumber, bool, str, bytes, np.ndarray, List[Any],


import os
import shutil
from absl import logging
from absl.testing import absltest
from envlogger.backends.python import riegeli_dataset_reader
from envlogger.backends.python import riegeli_dataset_writer
from envlogger.converters import codec
from envlogger.proto import storage_pb2
from google.protobuf import descriptor_pool
from google.protobuf import message_factory
from pybind11_abseil import status
  def test_clone(self):
    """Ensures that we can read the same data with a cloned reader."""
    writer = riegeli_dataset_writer.RiegeliDatasetWriter()
    try:
      writer.init(data_dir=self._directory)
    except status.StatusNotOk:
      logging.exception('Failed to initialize writer')

    for i in range(10):
      writer.add_step(codec.encode(i))
    writer.close()

    reader = riegeli_dataset_reader.RiegeliDatasetReader()
    try:
      reader.init(data_dir=self._directory)
    except status.StatusNotOk:
      logging.exception('Failed to initialize reader')

    cloned = reader.clone()
    self.assertEqual(cloned.num_steps, reader.num_steps)

    for i in range(reader.num_steps):
      step = reader.step(i)
      self.assertEqual(codec.decode(step), i)
    reader.close()

    # Even after closing the original `reader`, the cloned reader should still
    # work just like it.
    for i in range(cloned.num_steps):
      step = cloned.step(i)

# Path: envlogger/converters/codec.py
# _SCALAR_DIM_SIZE = -438
# def _python_int_to_bytes(py_int: int) -> bytes:
# def _set_datum_values_from_scalar(scalar: Union[ScalarNumber, bool, str, bytes],
#                                   datum: storage_pb2.Datum) -> bool:
# def _set_datum_values_from_array(array: np.ndarray,
#                                  values: storage_pb2.Datum.Values) -> None:
# def encode(
#     user_data: Union[np.ndarray, List[Any], Tuple[Any, ...], Dict[str, Any]]
# ) -> storage_pb2.Data:
# def decode_datum(
#     datum: storage_pb2.Datum
# ) -> Union[np.ndarray, ScalarNumber, bool, str, bytes]:
#     def from_bigint(int_bytes):
# def decode(
#     user_data: storage_pb2.Data
# ) -> Optional[Union[ScalarNumber, bool, str, bytes, np.ndarray, List[Any],


# Path: envlogger/backends/python/riegeli_dataset_test.py
import os
import shutil
from absl import logging
from absl.testing import absltest
from envlogger.backends.python import riegeli_dataset_reader
from envlogger.backends.python import riegeli_dataset_writer
from envlogger.converters import codec
from envlogger.proto import storage_pb2
from google.protobuf import descriptor_pool
from google.protobuf import message_factory
from pybind11_abseil import status

  def test_clone(self):
    """Ensures that we can read the same data with a cloned reader."""
    writer = riegeli_dataset_writer.RiegeliDatasetWriter()
    try:
      writer.init(data_dir=self._directory)
    except status.StatusNotOk:
      logging.exception('Failed to initialize writer')

    for i in range(10):
      writer.add_step(codec.encode(i))
    writer.close()

    reader = riegeli_dataset_reader.RiegeliDatasetReader()
    try:
      reader.init(data_dir=self._directory)
    except status.StatusNotOk:
      logging.exception('Failed to initialize reader')

    cloned = reader.clone()
    self.assertEqual(cloned.num_steps, reader.num_steps)

    for i in range(reader.num_steps):
      step = reader.step(i)
      self.assertEqual(codec.decode(step), i)
    reader.close()

    # Even after closing the original `reader`, the cloned reader should still
    # work just like it.
    for i in range(cloned.num_steps):
      step = cloned.step(i)

      self.assertEqual(codec.decode(step), i)