diff --git a/python/pandaset/sensors.py b/python/pandaset/sensors.py index a3dec3f..4a17df6 100644 --- a/python/pandaset/sensors.py +++ b/python/pandaset/sensors.py @@ -202,6 +202,16 @@ def __init__(self, directory: str) -> None: self._sensor_id = -1 Sensor.__init__(self, directory) + def _load_data_structure(self) -> None: + gz = sorted(glob.glob(f'{self._directory}/*.pkl.gz')) + plain = sorted(glob.glob(f'{self._directory}/*.pkl')) + if gz: + self._data_structure = gz + elif plain: + self._data_structure = plain + else: + super()._load_data_structure() + @overload def __getitem__(self, item: int) -> DataFrame: ... diff --git a/python/tests/test_lidar_extensions.py b/python/tests/test_lidar_extensions.py new file mode 100644 index 0000000..521a1e3 --- /dev/null +++ b/python/tests/test_lidar_extensions.py @@ -0,0 +1,40 @@ +"""Tests for LiDAR file extension fallback (#140).""" + +import tempfile +import unittest +from pathlib import Path + +import pandas as pd + +from pandaset.sensors import Lidar + + +class TestLidarExtensions(unittest.TestCase): + def test_loads_plain_pkl_when_gz_missing(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + pd.DataFrame({"x": [1.0], "y": [2.0]}).to_pickle(root / "00.pkl") + (root / "poses.json").write_text("[]") + (root / "timestamps.json").write_text("[]") + lidar = Lidar(str(root)) + lidar.load() + self.assertEqual(len(lidar.data), 1) + self.assertIn("x", lidar.data[0].columns) + + def test_prefers_pkl_gz_over_plain_pkl(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + pd.DataFrame({"x": [1.0]}).to_pickle( + root / "00.pkl.gz", compression="gzip" + ) + pd.DataFrame({"x": [999.0]}).to_pickle(root / "00.pkl") + (root / "poses.json").write_text("[]") + (root / "timestamps.json").write_text("[]") + lidar = Lidar(str(root)) + lidar.load() + self.assertEqual(len(lidar.data), 1) + self.assertEqual(lidar.data[0]["x"].iloc[0], 1.0) + + +if __name__ == "__main__": + unittest.main()