在 python 中创建包装上下文管理器到装饰器中的正确方法？

Question

我有几个网页想使用硒来抓取。我想自动执行此操作并在远程计算机上运行它。由于每个网站都不同，因此脚本需要不同的功能来完成工作。我没有让每个脚本都具有相同的代码来启动虚拟显示器和网络驱动程序，而是使用一个可以启动虚拟显示器和网络驱动程序的装饰器，如下所示：

    def open_headless_browser(func: Callable) -> Callable:
        disp = Display(visible=False, size=(100, 100))
        options = webdriver.ChromeOptions()
        options.add_argument("--headless=new")
        options.add_argument("--dns-prefetch-disable")
        def start(): -> None
            with disp as display:
                with webdriver.Chrome(options=self.options) as wd:
                    func()
        return start

然后我可以像这样拥有我的脚本（实际执行抓取的脚本）：

@open_headless_browser
def scrape_abc(url_abc: str) -> None:
    driver.get(url_abc)
    driver.find_elements_by_xpath('abc')

@open_headless_browser
def scrape_xyz(url_xyz: str) -> None:
    driver.get(url_xyz)
    driver.find_elements_by_css('xyz')

但是，有几件事与我有关：

• 我的

  scrape_abc

  和

  scrape_xzy

  函数中的代码有点尴尬，因为它不知道

  driver

  是什么（因为它是在装饰器中定义的）。
• 这还能用吗？我是否把事情过于复杂化了，或者我只是错误地理解了这个想法？
• 这是Pythonic吗

我在 python3.10 selenium4.15 pyvirtualdisplay3.0

编辑：经过一番思考，这种方法终究行不通。装饰函数将无法访问装饰器中定义的 webdriver 对象

Answer 1

编辑：经过一番思考，这种方法终究行不通。装饰函数将无法访问装饰器中定义的 webdriver 对象

当然可以，您只需将

wd

作为参数传递给函数，如下所示：

def open_headless_browser(func: Callable) -> Callable:
    disp = Display(visible=False, size=(100, 100))
    options = webdriver.ChromeOptions()
    options.add_argument("--headless=new")
    options.add_argument("--dns-prefetch-disable")
    def start(): -> None
        with disp as display:
            with webdriver.Chrome(options=options) as wd:
                func(wd)
    return start

那么你的函数将如下所示：

@open_headless_browser
def scrape_abc(driver: webdriver.Chrome) -> None:
    driver.get(url_abc)
    driver.find_elements_by_xpath('abc')

@open_headless_browser
def scrape_abc(driver: webdriver.Chrome) -> None:
    driver.get(url_xyz)
    driver.find_elements_by_xpath('xyz')

如果您希望能够传入 URL，则还需要在包装函数中定义参数：

def open_headless_browser(func: Callable) -> Callable:
    disp = Display(visible=False, size=(100, 100))
    options = webdriver.ChromeOptions()
    options.add_argument("--headless=new")
    options.add_argument("--dns-prefetch-disable")
    def start(url: str): -> None
        with disp as display:
            with webdriver.Chrome(options=options) as wd:
                func(wd, url)
    return start

@open_headless_browser
def scrape_abc(driver: webdriver.Chrome, url: str) -> None:
    driver.get(url)
    driver.find_elements_by_xpath('abc')

那么只需记住，虽然您将函数定义为具有两个参数，但您只用一个参数来调用它。