1 files changed, 157 insertions, 67 deletions

diff --git a/g4f/Provider/bing/create_images.py b/g4f/Provider/bing/create_images.py
index a1ecace3..060cd184 100644
--- a/g4f/Provider/bing/create_images.py
+++ b/g4f/Provider/bing/create_images.py
@@ -1,9 +1,16 @@
+"""
+This module provides functionalities for creating and managing images using Bing's service.
+It includes functions for user login, session creation, image creation, and processing.
+"""
+
 import asyncio
-import time, json, os
+import time
+import json
+import os
 from aiohttp import ClientSession
 from bs4 import BeautifulSoup
 from urllib.parse import quote
-from typing import Generator
+from typing import Generator, List, Dict
 
 from ..create_images import CreateImagesProvider
 from ..helper import get_cookies, get_event_loop
@@ -12,23 +19,47 @@ from ...base_provider import ProviderType
 from ...image import format_images_markdown
 
 BING_URL = "https://www.bing.com"
+TIMEOUT_LOGIN = 1200
+TIMEOUT_IMAGE_CREATION = 300
+ERRORS = [
+    "this prompt is being reviewed",
+    "this prompt has been blocked",
+    "we're working hard to offer image creator in more languages",
+    "we can't create your images right now"
+]
+BAD_IMAGES = [
+    "https://r.bing.com/rp/in-2zU3AJUdkgFe7ZKv19yPBHVs.png",
+    "https://r.bing.com/rp/TX9QuO3WzcCJz1uaaSwQAz39Kb0.jpg",
+]
+
+def wait_for_login(driver: WebDriver, timeout: int = TIMEOUT_LOGIN) -> None:
+    """
+    Waits for the user to log in within a given timeout period.
 
-def wait_for_login(driver: WebDriver, timeout: int = 1200) -> None:
+    Args:
+        driver (WebDriver): Webdriver for browser automation.
+        timeout (int): Maximum waiting time in seconds.
+
+    Raises:
+        RuntimeError: If the login process exceeds the timeout.
+    """
     driver.get(f"{BING_URL}/")
-    value = driver.get_cookie("_U")
-    if value:
-        return
     start_time = time.time()
-    while True:
+    while not driver.get_cookie("_U"):
         if time.time() - start_time > timeout:
             raise RuntimeError("Timeout error")
-        value = driver.get_cookie("_U")
-        if value:
-            time.sleep(1)
-            return
         time.sleep(0.5)
 
-def create_session(cookies: dict) -> ClientSession:
+def create_session(cookies: Dict[str, str]) -> ClientSession:
+    """
+    Creates a new client session with specified cookies and headers.
+
+    Args:
+        cookies (Dict[str, str]): Cookies to be used for the session.
+
+    Returns:
+        ClientSession: The created client session.
+    """
     headers = {
         "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
         "accept-encoding": "gzip, deflate, br",
@@ -47,28 +78,32 @@ def create_session(cookies: dict) -> ClientSession:
         "upgrade-insecure-requests": "1",
     }
     if cookies:
-        headers["cookie"] = "; ".join(f"{k}={v}" for k, v in cookies.items())
+        headers["Cookie"] = "; ".join(f"{k}={v}" for k, v in cookies.items())
     return ClientSession(headers=headers)
 
-async def create_images(session: ClientSession, prompt: str, proxy: str = None, timeout: int = 300) -> list:
-    url_encoded_prompt = quote(prompt)    
+async def create_images(session: ClientSession, prompt: str, proxy: str = None, timeout: int = TIMEOUT_IMAGE_CREATION) -> List[str]:
+    """
+    Creates images based on a given prompt using Bing's service.
+
+    Args:
+        session (ClientSession): Active client session.
+        prompt (str): Prompt to generate images.
+        proxy (str, optional): Proxy configuration.
+        timeout (int): Timeout for the request.
+
+    Returns:
+        List[str]: A list of URLs to the created images.
+
+    Raises:
+        RuntimeError: If image creation fails or times out.
+    """
+    url_encoded_prompt = quote(prompt)
     payload = f"q={url_encoded_prompt}&rt=4&FORM=GENCRE"
     url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=4&FORM=GENCRE"
-    async with session.post(
-        url,
-        allow_redirects=False,
-        data=payload,
-        timeout=timeout,
-    ) as response:
+    async with session.post(url, allow_redirects=False, data=payload, timeout=timeout) as response:
         response.raise_for_status()
-        errors = [
-            "this prompt is being reviewed",
-            "this prompt has been blocked",
-            "we're working hard to offer image creator in more languages",
-            "we can't create your images right now"
-        ]
         text = (await response.text()).lower()
-        for error in errors:
+        for error in ERRORS:
             if error in text:
                 raise RuntimeError(f"Create images failed: {error}")
     if response.status != 302:
@@ -107,54 +142,109 @@ async def create_images(session: ClientSession, prompt: str, proxy: str = None,
         raise RuntimeError(error)
     return read_images(text)
 
-def read_images(text: str) -> list:
-    html_soup = BeautifulSoup(text, "html.parser")
-    tags = html_soup.find_all("img")
-    image_links = [img["src"] for img in tags if "mimg" in img["class"]]
-    images = [link.split("?w=")[0] for link in image_links]
-    bad_images = [
-        "https://r.bing.com/rp/in-2zU3AJUdkgFe7ZKv19yPBHVs.png",
-        "https://r.bing.com/rp/TX9QuO3WzcCJz1uaaSwQAz39Kb0.jpg",
-    ]
-    if any(im in bad_images for im in images):
+def read_images(html_content: str) -> List[str]:
+    """
+    Extracts image URLs from the HTML content.
+
+    Args:
+        html_content (str): HTML content containing image URLs.
+
+    Returns:
+        List[str]: A list of image URLs.
+    """
+    soup = BeautifulSoup(html_content, "html.parser")
+    tags = soup.find_all("img", class_="mimg")
+    images = [img["src"].split("?w=")[0] for img in tags]
+    if any(im in BAD_IMAGES for im in images):
         raise RuntimeError("Bad images found")
     if not images:
         raise RuntimeError("No images found")
     return images
 
-async def create_images_markdown(cookies: dict, prompt: str, proxy: str = None) -> str:
-    session = create_session(cookies)
-    try:
+async def create_images_markdown(cookies: Dict[str, str], prompt: str, proxy: str = None) -> str:
+    """
+    Creates markdown formatted string with images based on the prompt.
+
+    Args:
+        cookies (Dict[str, str]): Cookies to be used for the session.
+        prompt (str): Prompt to generate images.
+        proxy (str, optional): Proxy configuration.
+
+    Returns:
+        str: Markdown formatted string with images.
+    """
+    async with create_session(cookies) as session:
         images = await create_images(session, prompt, proxy)
         return format_images_markdown(images, prompt)
-    finally:
-        await session.close()
 
-def get_cookies_from_browser(proxy: str = None) -> dict:
-    driver = get_browser(proxy=proxy)
-    try:
+def get_cookies_from_browser(proxy: str = None) -> Dict[str, str]:
+    """
+    Retrieves cookies from the browser using webdriver.
+
+    Args:
+        proxy (str, optional): Proxy configuration.
+
+    Returns:
+        Dict[str, str]: Retrieved cookies.
+    """
+    with get_browser(proxy=proxy) as driver:
         wait_for_login(driver)
+        time.sleep(1)
         return get_driver_cookies(driver)
-    finally:
-        driver.quit()
-
-def create_completion(prompt: str, cookies: dict = None, proxy: str = None) -> Generator:
-    loop = get_event_loop()
-    if not cookies:
-        cookies = get_cookies(".bing.com")
-    if "_U" not in cookies:
-        login_url = os.environ.get("G4F_LOGIN_URL")
-        if login_url:
-            yield f"Please login: [Bing]({login_url})\n\n"
-        cookies = get_cookies_from_browser(proxy)
-    yield loop.run_until_complete(create_images_markdown(cookies, prompt, proxy))
-
-async def create_async(prompt: str, cookies: dict = None, proxy: str = None) -> str:
-    if not cookies:
-        cookies = get_cookies(".bing.com")
-    if "_U" not in cookies:
-        cookies = get_cookies_from_browser(proxy)
-    return await create_images_markdown(cookies, prompt, proxy)
+
+class CreateImagesBing:
+    """A class for creating images using Bing."""
+
+    _cookies: Dict[str, str] = {}
+    
+    @classmethod
+    def create_completion(cls, prompt: str, cookies: Dict[str, str] = None, proxy: str = None) -> Generator[str, None, None]:
+        """
+        Generator for creating imagecompletion based on a prompt.
+
+        Args:
+            prompt (str): Prompt to generate images.
+            cookies (Dict[str, str], optional): Cookies for the session. If None, cookies are retrieved automatically.
+            proxy (str, optional): Proxy configuration.
+
+        Yields:
+            Generator[str, None, None]: The final output as markdown formatted string with images.
+        """
+        loop = get_event_loop()
+        cookies = cookies or cls._cookies or get_cookies(".bing.com")
+        if "_U" not in cookies:
+            login_url = os.environ.get("G4F_LOGIN_URL")
+            if login_url:
+                yield f"Please login: [Bing]({login_url})\n\n"
+            cls._cookies = cookies = get_cookies_from_browser(proxy)
+        yield loop.run_until_complete(create_images_markdown(cookies, prompt, proxy))
+
+    @classmethod
+    async def create_async(cls, prompt: str, cookies: Dict[str, str] = None, proxy: str = None) -> str:
+        """
+        Asynchronously creates a markdown formatted string with images based on the prompt.
+
+        Args:
+            prompt (str): Prompt to generate images.
+            cookies (Dict[str, str], optional): Cookies for the session. If None, cookies are retrieved automatically.
+            proxy (str, optional): Proxy configuration.
+
+        Returns:
+            str: Markdown formatted string with images.
+        """
+        cookies = cookies or cls._cookies or get_cookies(".bing.com")
+        if "_U" not in cookies:
+            cls._cookies = cookies = get_cookies_from_browser(proxy)
+        return await create_images_markdown(cookies, prompt, proxy)
 
 def patch_provider(provider: ProviderType) -> CreateImagesProvider:
-    return CreateImagesProvider(provider, create_completion, create_async)
\ No newline at end of file
+    """
+    Patches a provider to include image creation capabilities.
+
+    Args:
+        provider (ProviderType): The provider to be patched.
+
+    Returns:
+        CreateImagesProvider: The patched provider with image creation capabilities.
+    """
+    return CreateImagesProvider(provider, CreateImagesBing.create_completion, CreateImagesBing.create_async)
\ No newline at end of file