Skip to content

Commit e6b42ad

Browse files
committed
feat(client): implement context manager protocol
1 parent f4a89fe commit e6b42ad

File tree

3 files changed

+44
-2
lines changed

3 files changed

+44
-2
lines changed

python/rnet/__init__.pyi

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1354,6 +1354,11 @@ class Client:
13541354
asyncio.run(main())
13551355
```
13561356
"""
1357+
1358+
async def __aenter__(self) -> Any: ...
1359+
async def __aexit__(
1360+
self, _exc_type: Any, _exc_value: Any, _traceback: Any
1361+
) -> Any: ...
13571362

13581363
async def delete(
13591364
url: str,

python/rnet/blocking.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ def close(self) -> None:
127127
in future versions.
128128
"""
129129

130-
def __enter__(self) -> "Response": ...
130+
def __enter__(self) -> Any: ...
131131
def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: ...
132132
def __str__(self) -> str: ...
133133

@@ -204,7 +204,7 @@ def close(
204204
* `reason` - An optional reason for closing.
205205
"""
206206

207-
def __enter__(self) -> "WebSocket": ...
207+
def __enter__(self) -> Any: ...
208208
def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: ...
209209
def __str__(self) -> str: ...
210210

@@ -463,6 +463,9 @@ def get(
463463
```
464464
"""
465465
...
466+
467+
def __enter__(self) -> Any: ...
468+
def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: ...
466469

467470

468471
def delete(

src/client.rs

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -585,6 +585,21 @@ impl Client {
585585
}
586586
}
587587

588+
#[pymethods]
589+
impl Client {
590+
#[inline]
591+
async fn __aenter__(slf: Py<Self>) -> PyResult<Py<Self>> {
592+
Ok(slf)
593+
}
594+
595+
#[inline]
596+
async fn __aexit__(&self, _exc_type: Py<PyAny>, _exc_val: Py<PyAny>, _traceback: Py<PyAny>) {
597+
// TODO: Implement connection closing logic if necessary.
598+
}
599+
}
600+
601+
// ===== impl BlockingClient =====
602+
588603
#[pymethods]
589604
impl BlockingClient {
590605
/// Creates a new blocking Client instance.
@@ -729,3 +744,22 @@ impl BlockingClient {
729744
})
730745
}
731746
}
747+
748+
#[pymethods]
749+
impl BlockingClient {
750+
#[inline]
751+
fn __enter__(slf: PyRef<Self>) -> PyRef<Self> {
752+
slf
753+
}
754+
755+
#[inline]
756+
fn __exit__<'py>(
757+
&self,
758+
_py: Python<'py>,
759+
_exc_type: &Bound<'py, PyAny>,
760+
_exc_value: &Bound<'py, PyAny>,
761+
_traceback: &Bound<'py, PyAny>,
762+
) {
763+
// TODO: Implement connection closing logic if necessary.
764+
}
765+
}

0 commit comments

Comments
 (0)