From 90642d4d98d1317425f594879acbd030ad67a655 Mon Sep 17 00:00:00 2001
From: Jakub Ziebinski <ziebinskijakub@gmail.com>
Date: Wed, 19 Mar 2025 09:32:48 +0100
Subject: [PATCH] Implementation of the beekeepy objects to make it easier to
 create a beekeeper set

---
 beekeepy/beekeepy/objects/__init__.py     | 12 ++++++++
 beekeepy/beekeepy/objects/asynchronous.py | 36 +++++++++++++++++++++++
 beekeepy/beekeepy/objects/synchronous.py  | 36 +++++++++++++++++++++++
 3 files changed, 84 insertions(+)
 create mode 100644 beekeepy/beekeepy/objects/__init__.py
 create mode 100644 beekeepy/beekeepy/objects/asynchronous.py
 create mode 100644 beekeepy/beekeepy/objects/synchronous.py

diff --git a/beekeepy/beekeepy/objects/__init__.py b/beekeepy/beekeepy/objects/__init__.py
new file mode 100644
index 00000000..7a6390c3
--- /dev/null
+++ b/beekeepy/beekeepy/objects/__init__.py
@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from .asynchronous import BeekeepyObjects as AsyncBeekeepyObjects
+from .asynchronous import create_beekeeper_set as create_async_beekeeper_set
+from .synchronous import BeekeepyObjects, create_beekeeper_set
+
+__all__ = [
+    "AsyncBeekeepyObjects",
+    "create_async_beekeeper_set",
+    "BeekeepyObjects",
+    "create_beekeeper_set",
+]
diff --git a/beekeepy/beekeepy/objects/asynchronous.py b/beekeepy/beekeepy/objects/asynchronous.py
new file mode 100644
index 00000000..72d35d07
--- /dev/null
+++ b/beekeepy/beekeepy/objects/asynchronous.py
@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import AsyncGenerator
+
+from beekeepy import AsyncBeekeeper, AsyncSession, AsyncUnlockedWallet
+
+
+@dataclass
+class BeekeepyObjects:
+    wallet: AsyncUnlockedWallet
+    session: AsyncSession
+    beekeeper: AsyncBeekeeper
+
+
+@asynccontextmanager
+async def create_beekeeper_set(*, wallet_name: str, password: str) -> AsyncGenerator[BeekeepyObjects, None]:
+    """
+    Create a beekeeper set with a wallet and session.
+
+    This context manager handle the creation of a wallet and session and running beekeeper instance.
+    If you already have a wallet with the same name, it will unlock it.
+
+    Args:
+        wallet_name (str): The wallet to create/unlock name.
+        password (str): The wallet to create/unlock password.
+    """
+    async with await AsyncBeekeeper.factory() as beekeeper, await beekeeper.create_session() as session:
+        wallet = (
+            await session.create_wallet(name=wallet_name, password=password)
+            if wallet_name not in [w.name for w in await session.wallets_created]
+            else await (await session.open_wallet(name=wallet_name)).unlock(password=password)
+        )
+
+        yield BeekeepyObjects(wallet=wallet, session=session, beekeeper=beekeeper)
diff --git a/beekeepy/beekeepy/objects/synchronous.py b/beekeepy/beekeepy/objects/synchronous.py
new file mode 100644
index 00000000..ea1c634b
--- /dev/null
+++ b/beekeepy/beekeepy/objects/synchronous.py
@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Generator
+
+from beekeepy import Beekeeper, Session, UnlockedWallet
+
+
+@dataclass
+class BeekeepyObjects:
+    wallet: UnlockedWallet
+    session: Session
+    beekeeper: Beekeeper
+
+
+@contextmanager
+def create_beekeeper_set(*, wallet_name: str, password: str) -> Generator[BeekeepyObjects, None, None]:
+    """
+    Create a beekeeper set with a wallet and session.
+
+    This context manager handle the creation of a wallet and session and running beekeeper instance.
+    If you already have a wallet with the same name, it will unlock it.
+
+    Args:
+        wallet_name (str): The wallet to create/unlock name.
+        password (str): The wallet to create/unlock password.
+    """
+    with Beekeeper.factory() as beekeeper, beekeeper.create_session() as session:
+        wallet = (
+            session.create_wallet(name=wallet_name, password=password)
+            if wallet_name not in [w.name for w in session.wallets_created]
+            else (session.open_wallet(name=wallet_name)).unlock(password=password)
+        )
+
+        yield BeekeepyObjects(wallet=wallet, session=session, beekeeper=beekeeper)
-- 
GitLab