feat: wrap AI, Images, Analytics Engine, Vectorize and RateLimit Bindings to accept native Python objects

packages/cli/tests/bindings-test/src/test_analytics_engine.py

@@ -0,0 +1,23 @@
+import pytest
+from workers._workers import _BindingWrapper
+
+
+@pytest.mark.asyncio
+async def test_is_wrapped(env):
+    assert isinstance(env.ANALYTICS, _BindingWrapper)
+
+
+@pytest.mark.asyncio
+async def test_write_data_point_blobs_and_doubles(env):
+    env.ANALYTICS.writeDataPoint(
+        {
+            "blobs": ["blob1", "blob2"],
+            "doubles": [1.0, 2.5],
+            "indexes": ["idx"],
+        }
+    )
+
+
+@pytest.mark.asyncio
+async def test_write_data_point_empty(env):
+    env.ANALYTICS.writeDataPoint({})

packages/cli/tests/bindings-test/src/test_images.py

@@ -0,0 +1,40 @@
+import base64
+
+import js
+import pytest
+from pyodide.ffi import create_proxy, to_js
+from workers._workers import _BindingWrapper
+
+PNG_1X1 = base64.b64decode(
+    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
+)
+
+
+def _make_stream(data):
+    def start(controller):
+        controller.enqueue(to_js(data))
+        controller.close()
+
+    return js.ReadableStream.new(to_js({"start": create_proxy(start)}))
+
+
+@pytest.mark.asyncio
+async def test_is_wrapped(env):
+    assert isinstance(env.IMAGES, _BindingWrapper)
+
+
+@pytest.mark.asyncio
+async def test_output_as_png(env):
+    pipeline = env.IMAGES.input(_make_stream(PNG_1X1))
+    output = await pipeline.output({"format": "image/png"})
+    resp = output.response()
+    assert resp.headers.get("content-type") == "image/png"
+
+
+@pytest.mark.asyncio
+async def test_transform_and_output(env):
+    pipeline = env.IMAGES.input(_make_stream(PNG_1X1))
+    transformed = pipeline.transform({"width": 1, "height": 1})
+    output = await transformed.output({"format": "image/png"})
+    resp = output.response()
+    assert resp.headers.get("content-type") == "image/png"

packages/cli/tests/bindings-test/src/test_ratelimit.py

@@ -0,0 +1,13 @@
+import pytest
+from workers._workers import _BindingWrapper
+
+
+@pytest.mark.asyncio
+async def test_is_wrapped(env):
+    assert isinstance(env.RATE_LIMITER, _BindingWrapper)
+
+
+@pytest.mark.asyncio
+async def test_limit_success(env):
+    result = await env.RATE_LIMITER.limit({"key": "test-key"})
+    assert result.success is True

packages/cli/tests/bindings-test/src/test_vectorize.py

@@ -0,0 +1,10 @@
+import pytest
+from workers._workers import _BindingWrapper
+
+
+@pytest.mark.asyncio
+async def test_vectorize_is_wrapped(env):
+    # Vectorize requires remote database even in local development environment
+    # so we cannot test it in this unittest
+    # Here, we just make sure it is wrapped properly with our bindings wrapper
+    assert isinstance(env.VECTORIZE, _BindingWrapper)

packages/cli/tests/bindings-test/wrangler.jsonc

@@ -31,5 +31,15 @@
     },
     "migrations": [
         { "tag": "v1", "new_sqlite_classes": ["TestDurableObject"] }
+    ],
+    "vectorize": [
+        { "binding": "VECTORIZE", "index_name": "test-index" }
+    ],
+    "analytics_engine_datasets": [
+        { "binding": "ANALYTICS", "dataset": "test-dataset" }
+    ],
+    "images": { "binding": "IMAGES" },
+    "ratelimits": [
+        { "name": "RATE_LIMITER", "namespace_id": "1001", "simple": { "period": 60, "limit": 100 } }
     ]
 }

packages/runtime-sdk/src/workers/_workers.py

@@ -1313,6 +1313,13 @@ class _EnvWrapper:
         "R2Bucket",
         "D1Database",
         "WorkerQueue",
+        "Ai",
+        "VectorizeIndexImpl",
+        "AnalyticsEngineDataset",
+        "LocalAnalyticsEngineDataset",
+        "ImagesBindingImpl",
+        "HostedImagesBindingImpl",
+        "Ratelimit",
     }
 
     def __init__(self, env: Any):