Refactored the aiohttp_session parameter to use a decorator to access or create it when used as a function parameter name

surugaya_api/product.py

@@ -7,6 +7,30 @@ from bs4 import BeautifulSoup
 
 from surugaya_api.consts import SURU_COOKIE_STRING
 
+import inspect
+
+def get_or_create_aiohttp_session(func):
+    async def wrapper(*args, **kwargs):
+        sig = inspect.signature(func)
+        bound_args = sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+
+        aiohttp_session = bound_args.arguments.get("aiohttp_session")
+        is_created = False
+
+        if not aiohttp_session:
+            is_created = True
+            aiohttp_session = aiohttp.ClientSession()
+            bound_args.arguments["aiohttp_session"] = aiohttp_session
+
+        result = await func(*bound_args.arguments.values())
+
+        if is_created:
+            await aiohttp_session.close()
+
+        return result
+
+    return wrapper
 @dataclass
 class ProductStock:
     condition: str
@@ -34,30 +58,23 @@ class Product:
     id: int
     name: str
     main_image_href: str
-    stock: [ProductStock]
+    stock: list[ProductStock]
 
     @property
     def in_stock(self):
         return not self.stock
 
 
+@get_or_create_aiohttp_session
 async def load_product(product_id, aiohttp_session=None):
-    if not aiohttp_session:
+    async with aiohttp_session.get(
-        _aiohttp_session = aiohttp.ClientSession()
+        url="https://www.suruga-ya.jp/product/detail/" + str(product_id),
-    else:
+        headers={
-        _aiohttp_session = aiohttp_session
+            "Cookie": SURU_COOKIE_STRING,
-
+        },
-    async with _aiohttp_session.get(
-            url="https://www.suruga-ya.jp/product/detail/" + str(product_id),
-            headers={
-                "Cookie": SURU_COOKIE_STRING,
-            },
     ) as response:
         page = await response.text()
-
+        page_bs = BeautifulSoup(page, features="html.parser")
-        page_bs = BeautifulSoup(
-            page, features="html.parser"
-        )
 
         product = Product(
             id=int(product_id),
@@ -66,10 +83,9 @@ async def load_product(product_id, aiohttp_session=None):
             stock=[
                 ProductStock.from_item_price(item_price)
                 for item_price in page_bs.select(".item-price")
-            ]
+            ],
         )
 
-    if not aiohttp_session:
+    return product
-        await _aiohttp_session.close()
+
 
-    return product