From 38d92556d3c303d40eb24aaf0b7497faacedc20c Mon Sep 17 00:00:00 2001
From: TennyZhuang <zty0826@gmail.com>
Date: Tue, 24 Sep 2024 09:48:10 +0800
Subject: [PATCH] feat(bindings/python): export ConcurrentLimitLayer

---
 bindings/python/README.md                 |  2 +-
 bindings/python/python/opendal/layers.pyi |  4 ++++
 bindings/python/src/layers.rs             | 23 +++++++++++++++++++++++
 bindings/python/src/lib.rs                |  1 +
 bindings/python/tests/conftest.py         |  6 ++++--
 5 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/bindings/python/README.md b/bindings/python/README.md
index 2103a692452..38ddaab5464 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -41,7 +41,7 @@ async def main():
 asyncio.run(main())
 ```
 
-s3 service example: 
+s3 service example:
 ```python
 import opendal
 
diff --git a/bindings/python/python/opendal/layers.pyi b/bindings/python/python/opendal/layers.pyi
index ea25bdbd1e2..a4b2df62ea0 100644
--- a/bindings/python/python/opendal/layers.pyi
+++ b/bindings/python/python/opendal/layers.pyi
@@ -30,3 +30,7 @@ class RetryLayer(Layer):
         max_delay: Optional[float] = None,
         min_delay: Optional[float] = None,
     ) -> None: ...
+
+@final
+class ConcurrentLimitLayer(Layer):
+    def __init__(self, limit: int) -> None: ...
diff --git a/bindings/python/src/layers.rs b/bindings/python/src/layers.rs
index 4b78e27e43d..e03aafac9c9 100644
--- a/bindings/python/src/layers.rs
+++ b/bindings/python/src/layers.rs
@@ -80,3 +80,26 @@ impl RetryLayer {
         Ok(class)
     }
 }
+
+#[pyclass(module = "opendal.layers", extends=Layer)]
+#[derive(Clone)]
+pub struct ConcurrentLimitLayer(ocore::layers::ConcurrentLimitLayer);
+
+impl PythonLayer for ConcurrentLimitLayer {
+    fn layer(&self, op: Operator) -> Operator {
+        op.layer(self.0.clone())
+    }
+}
+
+#[pymethods]
+impl ConcurrentLimitLayer {
+    #[new]
+    #[pyo3(signature = (limit))]
+    fn new(limit: usize) -> PyResult<PyClassInitializer<Self>> {
+        let concurrent_limit = Self(ocore::layers::ConcurrentLimitLayer::new(limit));
+        let class = PyClassInitializer::from(Layer(Box::new(concurrent_limit.clone())))
+            .add_subclass(concurrent_limit);
+
+        Ok(class)
+    }
+}
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
index 1829892caf0..b499dc42aef 100644
--- a/bindings/python/src/lib.rs
+++ b/bindings/python/src/lib.rs
@@ -90,6 +90,7 @@ fn _opendal(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     let layers_module = PyModule::new_bound(py, "layers")?;
     layers_module.add_class::<Layer>()?;
     layers_module.add_class::<RetryLayer>()?;
+    layers_module.add_class::<ConcurrentLimitLayer>()?;
     m.add_submodule(&layers_module)?;
     py.import_bound("sys")?
         .getattr("modules")?
diff --git a/bindings/python/tests/conftest.py b/bindings/python/tests/conftest.py
index 86fa77f8a96..3c84ec1f353 100644
--- a/bindings/python/tests/conftest.py
+++ b/bindings/python/tests/conftest.py
@@ -60,8 +60,10 @@ def setup_config(service_name):
 
 @pytest.fixture(scope="session")
 def async_operator(service_name, setup_config):
-    return opendal.AsyncOperator(service_name, **setup_config).layer(
-        opendal.layers.RetryLayer()
+    return (
+        opendal.AsyncOperator(service_name, **setup_config)
+        .layer(opendal.layers.RetryLayer())
+        .layer(opendal.layers.ConcurrentLimitLayer(1024))
     )