Skip to content

Commit c391ea0

Browse files
committed
auto-detect thread count from argv
1 parent e57db51 commit c391ea0

File tree

2 files changed

+18
-4
lines changed

2 files changed

+18
-4
lines changed

autoparaselenium/__init__.py

+17-3
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import itertools as it
2+
import os
23
import sys
34
import time
45
from concurrent.futures import ThreadPoolExecutor
@@ -16,7 +17,7 @@
1617
_test_count = 0 # manual reference counting since threading borks with destructors
1718

1819

19-
def configure(*_, extensions: List[Extension] = [], headless=True, selenium_dir="drivers", threads=1):
20+
def configure(*_, extensions: List[Extension] = [], headless=True, selenium_dir="drivers"):
2021
global _browser_pool
2122

2223
if _browser_pool is not None:
@@ -26,7 +27,7 @@ def configure(*_, extensions: List[Extension] = [], headless=True, selenium_dir=
2627
for browser in [chrome, firefox]:
2728
browser.setup_driver(conf.selenium_dir)
2829

29-
_browser_pool = BrowserPool(conf, threads)
30+
_browser_pool = BrowserPool(conf, __get_threads())
3031

3132

3233
def run_on(*browsers):
@@ -47,7 +48,7 @@ def wrapper(test):
4748
to_run = [*map(partial(__wrap_test, test=test), browsers)]
4849

4950
def inner():
50-
if 1 or "--tests-per-worker" in sys.argv:
51+
if "--tests-per-worker" in sys.argv:
5152
with ThreadPoolExecutor(max_workers=len(to_run)) as pool:
5253
pool.map(lambda test_: test_(), to_run)
5354
else:
@@ -90,3 +91,16 @@ def inner():
9091
inner.__doc__ = test.__doc__
9192

9293
return inner
94+
95+
96+
def __get_threads(args=sys.argv):
97+
if "--tests-per-worker" not in args:
98+
return 1
99+
tests_per_worker_idx = args.index("--tests-per-worker")
100+
next_arg = "".join(args[tests_per_worker_idx + 1: tests_per_worker_idx + 2])
101+
if next_arg == "auto":
102+
return os.cpu_count() // 2 + 1
103+
try:
104+
return int("0" + next_arg) // 2 + 1
105+
except ValueError:
106+
return 1

tests/test_everything.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from autoparaselenium import configure, chrome, firefox, run_on
66

77

8-
configure(headless=False, threads=2)
8+
configure(headless=False)
99

1010
def log(*args):
1111
with open("log", "a+") as fout:

0 commit comments

Comments
 (0)