datago 2025.6.2__tar.gz → 2025.6.4__tar.gz

{datago-2025.6.2 → datago-2025.6.4}/Cargo.lock

@@ -500,7 +500,7 @@ dependencies = [
 
 [[package]]
 name = "datago"
-version = "2025.6.2"
+version = "2025.6.4"
 dependencies = [
  "async-compression",
  "async-tar",
@@ -525,6 +525,7 @@ dependencies = [
  "serde",
  "serde_json",
  "tar",
+ "tempfile",
  "threadpool",
  "tokio",
  "tokio-util",
@@ -1771,7 +1772,7 @@ checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94"
 dependencies = [
  "rand_chacha 0.9.0",
  "rand_core 0.9.3",
- "zerocopy 0.8.23",
+ "zerocopy 0.8.25",
 ]
 
 [[package]]
@@ -2981,11 +2982,11 @@ dependencies = [
 
 [[package]]
 name = "zerocopy"
-version = "0.8.23"
+version = "0.8.25"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fd97444d05a4328b90e75e503a34bad781f14e28a823ad3557f0750df1ebcbc6"
+checksum = "a1702d9583232ddb9174e01bb7c15a2ab8fb1bc6f227aa1233858c351a3ba0cb"
 dependencies = [
- "zerocopy-derive 0.8.23",
+ "zerocopy-derive 0.8.25",
 ]
 
 [[package]]
@@ -3001,9 +3002,9 @@ dependencies = [
 
 [[package]]
 name = "zerocopy-derive"
-version = "0.8.23"
+version = "0.8.25"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6352c01d0edd5db859a63e2605f4ea3183ddbd15e2c4a9e7d32184df75e4f154"
+checksum = "28a6e20d751156648aa063f3800b706ee209a32c0b4d9f24be3d980b01be55ef"
 dependencies = [
  "proc-macro2",
  "quote",

{datago-2025.6.2 → datago-2025.6.4}/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "datago"
 edition = "2021"
-version = "2025.6.2"
+version = "2025.6.4"
 
 [lib]
 # exposed by pyo3
@@ -43,8 +43,23 @@ bracoxide = "0.1.6"
 infer = "0.16.0"
 fast_image_resize = { version ="5.1.3", features=["image"]}
 
+[dev-dependencies]
+tempfile = "3.13.0"
+
 [profile.release]
 opt-level = 3
 lto = "fat"
 codegen-units = 1
 panic = "abort"
+
+[target.x86_64-apple-darwin]
+rustflags = [
+  "-C", "link-arg=-undefined",
+  "-C", "link-arg=dynamic_lookup",
+]
+
+[target.aarch64-apple-darwin]
+rustflags = [
+  "-C", "link-arg=-undefined",
+  "-C", "link-arg=dynamic_lookup",
+]

{datago-2025.6.2 → datago-2025.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datago
-Version: 2025.6.2
+Version: 2025.6.4
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

datago-2025.6.4/python/test_datago_client.py

@@ -0,0 +1,444 @@
+import json
+import tempfile
+import os
+from datago import DatagoClient, initialize_logging
+from PIL import Image
+
+
+def create_test_images(directory, count=5):
+    """Helper function to create test images in a directory."""
+    image_paths = []
+    for i in range(count):
+        img = Image.new(
+            "RGB", (100, 100), color=(i * 50 % 255, (i * 100) % 255, (i * 150) % 255)
+        )
+        path = os.path.join(directory, f"test_image_{i}.png")
+        img.save(path)
+        image_paths.append(path)
+    return image_paths
+
+
+class TestDatagoClient:
+    """Test cases for DatagoClient functionality."""
+
+    def test_initialize_logging(self):
+        """Test the initialize_logging function."""
+        # Should return True on first call
+        result = initialize_logging("info")
+        assert isinstance(result, bool)
+
+        # Test with None parameter
+        result = initialize_logging(None)
+        assert isinstance(result, bool)
+
+    def test_client_instantiation_file_source(self):
+        """Test creating a client with file source configuration."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            create_test_images(tmpdir, 3)
+
+            config = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": False,
+                },
+                "limit": 3,
+                "samples_buffer_size": 10,
+            }
+
+            client = DatagoClient(json.dumps(config))
+            assert client is not None
+
+    # We panic out at the moment if the config is invalid, so this test is commented out.
+    # Uncomment this test if you want to handle invalid configurations gracefully.
+    # def test_client_instantiation_invalid_config(self):
+    #     """Test that invalid configuration raises an error."""
+    #     invalid_config = '{"invalid": "config"}'
+
+    #     with pytest.raises((ValueError, RuntimeError)):
+    #         DatagoClient(invalid_config)
+
+    def test_client_start_stop_file_source(self):
+        """Test starting and stopping client with file source."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            create_test_images(tmpdir, 3)
+
+            config = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": False,
+                },
+                "limit": 3,
+                "samples_buffer_size": 10,
+            }
+
+            client = DatagoClient(json.dumps(config))
+            client.start()
+            client.stop()
+
+    def test_get_sample_file_source(self):
+        """Test getting samples from file source."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            create_test_images(tmpdir, 5)
+
+            config = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": False,
+                },
+                "limit": 3,
+                "samples_buffer_size": 10,
+            }
+
+            client = DatagoClient(json.dumps(config))
+
+            samples_received = []
+            for _ in range(3):
+                sample = client.get_sample()
+                if sample:
+                    samples_received.append(sample)
+                else:
+                    break
+
+            assert len(samples_received) <= 3
+
+            for sample in samples_received:
+                assert sample.id != ""
+                assert sample.source == "filesystem"
+                assert sample.image.width > 0
+                assert sample.image.height > 0
+                assert len(sample.image.data) > 0
+
+    def test_client_with_image_transformations(self):
+        """Test client with image transformation configuration."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            create_test_images(tmpdir, 3)
+
+            config = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": False,
+                },
+                "image_config": {
+                    "crop_and_resize": True,
+                    "default_image_size": 64,
+                    "downsampling_ratio": 16,
+                    "min_aspect_ratio": 0.5,
+                    "max_aspect_ratio": 2.0,
+                    "pre_encode_images": False,
+                    "image_to_rgb8": True,
+                },
+                "limit": 2,
+                "samples_buffer_size": 10,
+            }
+
+            client = DatagoClient(json.dumps(config))
+            sample = client.get_sample()
+
+            assert sample is not None
+            assert sample.image.width <= 64
+            assert sample.image.height <= 64
+            assert sample.image.channels == 3  # RGB8
+
+    def test_client_with_image_encoding(self):
+        """Test client with image encoding enabled."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            create_test_images(tmpdir, 3)
+
+            config = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": False,
+                },
+                "image_config": {
+                    "crop_and_resize": False,
+                    "pre_encode_images": True,
+                    "image_to_rgb8": False,
+                },
+                "limit": 2,
+                "samples_buffer_size": 10,
+            }
+
+            client = DatagoClient(json.dumps(config))
+            sample = client.get_sample()
+
+            assert sample is not None
+            assert sample.image.channels == -1  # Encoded images have channels = -1
+            assert len(sample.image.data) > 0
+
+    def test_random_sampling(self):
+        """Test that random sampling produces different results."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            create_test_images(tmpdir, 10)
+
+            config_base = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": True,
+                },
+                "limit": 5,
+                "samples_buffer_size": 10,
+            }
+
+            # Get two sets of samples with random sampling
+            client1 = DatagoClient(json.dumps(config_base))
+            samples1 = []
+            for _ in range(5):
+                sample = client1.get_sample()
+                if sample:
+                    samples1.append(sample.id)
+                else:
+                    break
+
+            client2 = DatagoClient(json.dumps(config_base))
+            samples2 = []
+            for _ in range(5):
+                sample = client2.get_sample()
+                if sample:
+                    samples2.append(sample.id)
+                else:
+                    break
+
+            # With random sampling, the samples should be different
+            assert len(samples1) > 0
+            assert len(samples2) > 0
+            assert (
+                len(set(samples1) & set(samples2)) < 5
+            )  # Expect some overlap, but not all
+
+    def test_world_size_and_rank(self):
+        """Test that different ranks get different subsets of data."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            create_test_images(tmpdir, 10)
+
+            config_rank0 = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 2,
+                    "random_sampling": False,
+                },
+                "limit": 10,
+                "samples_buffer_size": 10,
+            }
+
+            config_rank1 = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 1,
+                    "world_size": 2,
+                    "random_sampling": False,
+                },
+                "limit": 10,
+                "samples_buffer_size": 10,
+            }
+
+            client0 = DatagoClient(json.dumps(config_rank0))
+            samples0 = []
+            for _ in range(10):
+                sample = client0.get_sample()
+                if sample:
+                    samples0.append(sample.id)
+                else:
+                    break
+
+            client1 = DatagoClient(json.dumps(config_rank1))
+            samples1 = []
+            for _ in range(10):
+                sample = client1.get_sample()
+                if sample:
+                    samples1.append(sample.id)
+                else:
+                    break
+
+            # Different ranks should get different samples
+            assert len(samples0) > 0
+            assert len(samples1) > 0
+
+            # No overlap between ranks
+            overlap = set(samples0) & set(samples1)
+            assert len(overlap) == 0
+
+    def test_limit_respected(self):
+        """Test that the client respects the limit parameter."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            create_test_images(tmpdir, 10)
+
+            config = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": False,
+                },
+                "limit": 3,
+                "samples_buffer_size": 10,
+            }
+
+            client = DatagoClient(json.dumps(config))
+            samples_received = 0
+
+            while True:
+                sample = client.get_sample()
+                if sample:
+                    samples_received += 1
+                else:
+                    break
+
+                # Safety valve to prevent infinite loop
+                if samples_received > 10:
+                    break
+
+            # Should respect the limit (might have small buffer)
+            assert samples_received <= 4  # Allow small buffer
+
+    def test_empty_directory(self):
+        """Test client behavior with empty directory."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": False,
+                },
+                "limit": 3,
+                "samples_buffer_size": 10,
+            }
+
+            client = DatagoClient(json.dumps(config))
+            sample = client.get_sample()
+
+            # Should return None when no files available
+            assert sample is None
+
+    def test_nonexistent_directory(self):
+        """Test client behavior with nonexistent directory."""
+        config = {
+            "source_type": "file",
+            "source_config": {
+                "root_path": "/nonexistent/directory",
+                "rank": 0,
+                "world_size": 1,
+                "random_sampling": False,
+            },
+            "limit": 3,
+            "samples_buffer_size": 10,
+        }
+
+        client = DatagoClient(json.dumps(config))
+        sample = client.get_sample()
+
+        # Should handle gracefully and return None
+        assert sample is None
+
+    def test_client_drop_cleanup(self):
+        """Test that client cleans up properly when dropped."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            create_test_images(tmpdir, 3)
+
+            config = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": False,
+                },
+                "limit": 3,
+                "samples_buffer_size": 10,
+            }
+
+            client = DatagoClient(json.dumps(config))
+            client.start()
+
+            # Client should clean up when it goes out of scope
+            del client
+
+    def test_multiple_starts_stops(self):
+        """Test that multiple start/stop calls don't cause issues."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            create_test_images(tmpdir, 3)
+
+            config = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": False,
+                },
+                "limit": 3,
+                "samples_buffer_size": 10,
+            }
+
+            client = DatagoClient(json.dumps(config))
+
+            # Multiple starts should be safe
+            client.start()
+            client.start()
+
+            # Multiple stops should be safe
+            client.stop()
+            client.stop()
+
+    def test_various_image_formats(self):
+        """Test client with various image formats."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create images with different formats
+            formats = [
+                ("test1.png", "PNG"),
+                ("test2.jpg", "JPEG"),
+            ]
+
+            for filename, format_name in formats:
+                img = Image.new("RGB", (50, 50), color="red")
+                path = os.path.join(tmpdir, filename)
+                img.save(path, format=format_name)
+
+            config = {
+                "source_type": "file",
+                "source_config": {
+                    "root_path": tmpdir,
+                    "rank": 0,
+                    "world_size": 1,
+                    "random_sampling": False,
+                },
+                "limit": 4,
+                "samples_buffer_size": 10,
+            }
+
+            client = DatagoClient(json.dumps(config))
+
+            samples_received = 0
+            while True:
+                sample = client.get_sample()
+                if sample:
+                    samples_received += 1
+                    assert sample.image.width == 50
+                    assert sample.image.height == 50
+                else:
+                    break
+
+            assert samples_received == len(formats)