Skip to content

Commit

Permalink
formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
r2dev2 committed Jul 28, 2023
1 parent 69e394c commit 07097e1
Show file tree
Hide file tree
Showing 5 changed files with 31 additions and 22 deletions.
23 changes: 13 additions & 10 deletions autoparaselenium/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,18 +7,19 @@
from functools import partial, wraps
from typing import Iterable, List, Optional

from autoparaselenium.browsers import chrome, firefox
from autoparaselenium.browser_pool import BrowserPool
from autoparaselenium.browsers import chrome, firefox
from autoparaselenium.models import Conf, Extension


_browser_pool: Optional[BrowserPool] = None
all_ = [chrome, firefox]

_test_count = 0 # manual reference counting since threading borks with destructors
_test_count = 0 # manual reference counting since threading borks with destructors


def configure(*_, extensions: List[Extension] = [], headless=True, selenium_dir="drivers"):
def configure(
*_, extensions: List[Extension] = [], headless=True, selenium_dir="drivers"
):
global _browser_pool

if _browser_pool is not None:
Expand All @@ -37,7 +38,6 @@ def run_on(*browsers):
if not browsers:
raise TypeError("Please specify a browser or browser list to run on")


if isinstance(browsers[0], Iterable):
browsers = [*it.chain(*browsers)]

Expand Down Expand Up @@ -70,25 +70,28 @@ def __wrap_test(browser, test):
_test_count += 1

if _browser_pool is None:
raise RuntimeError("Please call autoparaselenium.configure() before creating tests")
raise RuntimeError(
"Please call autoparaselenium.configure() before creating tests"
)

def inner():
global _test_count

try:
_test_count -= 1
driver = _browser_pool.acquire(browser)
driver.get("data:,") # initialize driver website
driver.get("data:,") # initialize driver website
test(driver)
finally:
with suppress(Exception):
_browser_pool.release(driver)

if _test_count == 0:
time.sleep(0.10) # idk but seems like it needs a bit of time before you can close the pool
time.sleep(
0.10
) # idk but seems like it needs a bit of time before you can close the pool
_browser_pool.clean_up()


inner.__name__ = f"{test.__name__}__{'chrome' if browser is chrome else 'firefox'}"
inner.__doc__ = test.__doc__

Expand All @@ -99,7 +102,7 @@ def __get_threads(args=sys.argv):
if "--tests-per-worker" not in args:
return 1
tests_per_worker_idx = args.index("--tests-per-worker")
next_arg = "".join(args[tests_per_worker_idx + 1: tests_per_worker_idx + 2])
next_arg = "".join(args[tests_per_worker_idx + 1 : tests_per_worker_idx + 2])
if next_arg == "auto":
return os.cpu_count() // 2 + 1
try:
Expand Down
10 changes: 3 additions & 7 deletions autoparaselenium/browser_pool.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
from concurrent.futures import ThreadPoolExecutor
from contextlib import suppress
from threading import Lock
from queue import Queue
from threading import Lock

from selenium import webdriver

from autoparaselenium.models import Conf
from autoparaselenium.browsers import chrome, firefox
from autoparaselenium.models import Conf


class BrowserPool:
def __init__(self, conf: Conf, threads: int):
Expand All @@ -23,27 +24,22 @@ def __init__(self, conf: Conf, threads: int):
for driver in pool.map(self.__open_browser, [firefox] * threads):
self._firefoxes.put(driver)


def acquire(self, browser):
return self.__get_queue(browser).get(browser)


def release(self, driver):
browser = chrome if isinstance(driver, webdriver.Chrome) else firefox
self.__get_queue(browser).put(driver)


def __get_queue(self, browser):
return self._chromes if browser is chrome else self._firefoxes


def __open_browser(self, browser):
driver = browser.get_selenium(self._conf.selenium_dir, self._conf)
with suppress(Exception):
driver.switch_to.window("1")
return driver


def clean_up(self):
for q in [self._chromes, self._firefoxes]:
while not q.empty():
Expand Down
5 changes: 3 additions & 2 deletions autoparaselenium/browsers/chrome.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,8 +93,9 @@ def setup_driver(pwd: Path) -> None:
su.unzip,
]
for entry in (
r.json()["milestones"][str(chrome_version)]
["downloads"]["chromedriver"]
r.json()["milestones"][str(chrome_version)]["downloads"][
"chromedriver"
]
)
}
if (pwd / "chromedriver").exists():
Expand Down
11 changes: 9 additions & 2 deletions autoparaselenium/models.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,21 @@
from pathlib import Path
from typing import List, Optional


class Extension:
def __init__(self, *_, firefox: Optional[str]=None, chrome: Optional[str]=None):
def __init__(self, *_, firefox: Optional[str] = None, chrome: Optional[str] = None):
self.firefox = firefox
self.chrome = chrome


class Conf:
def __init__(self, *_, extensions: List[Extension] = [], headless=True, selenium_dir="drivers"):
def __init__(
self,
*_,
extensions: List[Extension] = [],
headless=True,
selenium_dir="drivers",
):
self.extensions = extensions
self.headless = headless
self.selenium_dir = Path(selenium_dir)
4 changes: 3 additions & 1 deletion autoparaselenium/setup_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,8 @@ def setup_driver(platform_install, platform_drivers, pwd: Path) -> None:
)


def __download_and_extract(url: str, extract: Callable[[str, str], None], pwd: Path) -> None:
def __download_and_extract(
url: str, extract: Callable[[str, str], None], pwd: Path
) -> None:
download(url, pwd / "temp")
extract(pwd / "temp", pwd)

0 comments on commit 07097e1

Please sign in to comment.