Make the subprocess browser class re-useable

This commit is contained in:
Kovid Goyal 2024-08-15 17:09:57 +05:30
parent 3e7bce7131
commit 03851ec6b3
No known key found for this signature in database
GPG Key ID: 06BC317B515ACE7C

View File

@ -4,6 +4,7 @@
import json import json
import os import os
import shutil import shutil
import subprocess
import time import time
from contextlib import suppress from contextlib import suppress
from io import BytesIO from io import BytesIO
@ -73,7 +74,7 @@ class Browser:
if start_worker: if start_worker:
self._ensure_state() self._ensure_state()
def open(self, url_or_request: Request, data=None, timeout=None): def _open(self, url_or_request: Request, data=None, timeout=None, visit: bool = True):
method = 'POST' if data else 'GET' method = 'POST' if data else 'GET'
headers = [] headers = []
if hasattr(url_or_request, 'get_method'): if hasattr(url_or_request, 'get_method'):
@ -100,12 +101,15 @@ class Browser:
if not has_header('Content-Type'): if not has_header('Content-Type'):
headers.append(('Content-Type', 'text/plain')) headers.append(('Content-Type', 'text/plain'))
if not self.is_method_ok(method):
raise KeyError(f'The HTTP {method} request method is not supported')
with self.lock: with self.lock:
self._ensure_state() self._ensure_state()
self.id_counter += 1 self.id_counter += 1
cmd = { cmd = {
'action': 'download', 'id': self.id_counter, 'url': url, 'method': method, 'timeout': timeout, 'action': 'download', 'id': self.id_counter, 'url': url, 'method': method, 'timeout': timeout,
'headers': self.addheaders + headers} 'headers': self.addheaders + headers, 'visit': visit,}
if data: if data:
with open(os.path.join(self.tdir, f'i{self.id_counter}'), 'wb') as f: with open(os.path.join(self.tdir, f'i{self.id_counter}'), 'wb') as f:
if hasattr(data, 'read'): if hasattr(data, 'read'):
@ -118,7 +122,14 @@ class Browser:
self._send_command(cmd) self._send_command(cmd)
return res return res
open_novisit = open def open(self, url_or_request: Request, data=None, timeout=None):
return self._open(url_or_request, data, timeout)
def open_novisit(self, url_or_request: Request, data=None, timeout=None):
return self._open(url_or_request, data, timeout, visit=False)
def is_method_ok(self, method: str) -> bool:
return True
def set_simple_cookie(self, name: str, value: str, domain: str | None = None, path: str | None = '/'): def set_simple_cookie(self, name: str, value: str, domain: str | None = None, path: str | None = '/'):
''' '''
@ -148,10 +159,13 @@ class Browser:
with self.lock: with self.lock:
if not self.tdir: if not self.tdir:
self.tdir = PersistentTemporaryDirectory() self.tdir = PersistentTemporaryDirectory()
self.worker = run_worker(self.tdir, self.user_agent, self.verify_ssl_certificates) self.worker = self.run_worker()
self.dispatcher = Thread(target=self._dispatch, daemon=True) self.dispatcher = Thread(target=self._dispatch, daemon=True)
self.dispatcher.start() self.dispatcher.start()
def run_worker(self) -> subprocess.Popen:
return run_worker(self.tdir, self.user_agent, self.verify_ssl_certificates)
def _dispatch(self): def _dispatch(self):
try: try:
for line in self.worker.stdout: for line in self.worker.stdout: