96 lines
2.1 KiB
Python
96 lines
2.1 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
import io
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
from contextlib import contextmanager
|
|
import shutil
|
|
|
|
from PIL import Image
|
|
import click
|
|
from marionette_driver.marionette import Marionette
|
|
from marionette_driver import expected, Wait, By
|
|
|
|
|
|
@contextmanager
|
|
def tempdir(*args, **kwargs):
|
|
path = tempfile.mkdtemp(*args, **kwargs)
|
|
try:
|
|
yield path
|
|
finally:
|
|
shutil.rmtree(path)
|
|
|
|
|
|
@contextmanager
|
|
def process(args):
|
|
process = subprocess.Popen(args)
|
|
try:
|
|
yield process
|
|
finally:
|
|
process.kill()
|
|
|
|
|
|
class Environ(object):
|
|
def __getattr__(self, name):
|
|
try:
|
|
return os.environ[name]
|
|
except KeyError:
|
|
raise AttributeError(name)
|
|
|
|
|
|
class MarionetteWrapper(Marionette):
|
|
def full_page_screenshot(self, filename):
|
|
img = Image.open(io.BytesIO(self.screenshot(element=None, format='binary', full=True)))
|
|
img.save(filename)
|
|
|
|
@property
|
|
def wait_until_expected(self):
|
|
return WaitUntil(self)
|
|
|
|
@property
|
|
def env(self):
|
|
return Environ()
|
|
|
|
|
|
class WaitUntilExpected:
|
|
def __init__(self, client, name):
|
|
self.client = client
|
|
self.name = name
|
|
|
|
def by_name(self, name):
|
|
return self(By.NAME, name)
|
|
|
|
def by_css_selector(self, selector):
|
|
return self(By.CSS_SELECTOR, selector)
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return Wait(self.client).until(getattr(expected, self.name)(*args, **kwargs))
|
|
|
|
|
|
class WaitUntil:
|
|
def __init__(self, client):
|
|
self.client = client
|
|
|
|
def __getattr__(self, name):
|
|
return WaitUntilExpected(self.client, name)
|
|
|
|
|
|
@click.command()
|
|
@click.argument('script_py')
|
|
def screenshot(script_py):
|
|
global client
|
|
|
|
with tempdir() as profile_path:
|
|
with process(['firefox', '--marionette', '--no-remote', '--new-instance', '--profile', profile_path, '--headless']):
|
|
client = MarionetteWrapper('localhost', '2828')
|
|
client.start_session()
|
|
with open(script_py) as fd:
|
|
exec(fd.read(), {})
|
|
client.cleanup()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
screenshot()
|