add 'tabs' parameter

deedy5 · Feb 19, 2024 · 8475838 · 8475838
1 parent a3125c3
commit 8475838
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/fake_traffic/cli.py b/fake_traffic/cli.py
@@ -37,6 +37,14 @@
     help="Run the browser in non-headless mode",
     required=False,
 )
+parser.add_argument(
+    "-t",
+    "--tabs",
+    default=5,
+    type=int,
+    help="Limit the number of tabs in browser. Defaults to 5",
+    required=False,
+)
 parser.add_argument(
     "-lf",
     "--logfile",
@@ -61,7 +69,7 @@
 language_split = args.language.split("-")
 language = f"{language_split[0]}-{language_split[1].upper()}"
 logging.info(
-    f"Run crawl with: {country=}, {language=}, category={args.category}, headless={args.headless}, logfile={args.logfile}"
+    f"Run crawl with: {country=}, {language=}, category={args.category}, headless={args.headless}, tabs={args.tabs}, logfile={args.logfile}"
 )
 
 
@@ -70,5 +78,6 @@
     language=language,
     category=args.category,
     headless=args.headless,
+    tabs=args.tabs,
 )
 fake_traffic.crawl()
diff --git a/fake_traffic/fake_traffic.py b/fake_traffic/fake_traffic.py
@@ -6,8 +6,6 @@
 from playwright_stealth import stealth_async
 
 logger = logging.getLogger("__name__")
-SEMAPHORE = asyncio.Semaphore(5)
-
 
 # playwright install chromium
 res = subprocess.run(
@@ -27,23 +25,27 @@ def __init__(
         language="en-US",
         category="h",
         headless=True,
+        tabs=5,
+
     ):
         """Internet traffic generator. Utilizes real-time google search trends by specified parameters.
         country = country code ISO 3166-1 Alpha-2 code (https://www.iso.org/obp/ui/),
         language = country-language code ISO-639 and ISO-3166 (https://www.fincher.org/Utilities/CountryLanguageList.shtml),
         category = category of interest of a user (defaults to 'h'):
                 'all' (all), 'b' (business), 'e' (entertainment),
                 'm' (health), 's' (sports), 't' (sci/tech), 'h' (top stories);
-        headless = True/False (defaults to True).
+        headless = True/False (defaults to True);
+        tabs = limit the number of tabs in browser (defaults to 5).
         """
         self.country = country
         self.language = language
         self.category = category
         self.headless = headless
         self.browser = None
+        self.semaphore = asyncio.Semaphore(tabs)
 
     async def abrowse(self, url):
-        async with SEMAPHORE:
+        async with self.semaphore:
             page = await self.browser.new_page()
             await stealth_async(page)
             try:
@@ -128,5 +130,6 @@ def crawl(self):
         language="en-US",
         category="h",
         headless=True,
+        tabs=5,
     )
     fake_traffic.crawl()