twitter_scraper/code/selenium_scraper.py

from selenium import webdriver
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
from selenium.webdriver.common.by import By
from selenium.webdriver.support.wait import WebDriverWait
import selenium.webdriver.support.expected_conditions as EC
from selenium.webdriver.remote.webelement import WebElement
from pydantic import BaseModel, validator, parse_obj_as, AnyHttpUrl
from typing import Union, List
from datetime import datetime, timedelta
from PIL import Image
from io import BytesIO
import requests
import pyperclip
import logging
import time
import traceback

from database import connect_mongo
from classes import Tweet

class Twitter(object):
    def __init__(self):
        self.webdriver = None
        self.actiondriver = None
        self.timeline_elem = None
        self._read_posts_timestamp_list = list()
        self._unread_posts_element_list = list()

    @staticmethod
    def _url( account: str, date_begin: datetime, date_end: datetime) -> AnyHttpUrl:
        return parse_obj_as(AnyHttpUrl, f"https://twitter.com/search?q=(from%3A{account})%20until%3A{date_end.strftime('%Y-%m-%d')}%20since%3A{date_begin.strftime('%Y-%m-%d')}%20-filter%3Areplies&src=typed_query&f=top")

    def setup(self) -> None:
        # setup webdriver
        driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()))
        self.webdriver = driver
        self.actiondriver = ActionChains(self.webdriver)

    def close(self) -> None:
        self.webdriver.close()

    def __enter__(self):
        self.setup()
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def _update_unread_posts(self) -> None:
        self._handle_popup()
        elem_list = self.timeline_elem.find_elements(By.XPATH, """.//article[@data-testid="tweet"]""")
        # elem_list = self.webdriver.find_elements(By.CSS_SELECTOR, '[data-testid="tweet"]')
        for elem in elem_list:
            try:
                elem_list = elem.find_elements(By.XPATH, './/time')
                assert len(elem_list) > 0, 'failed locating time of post'
                timestamp_str = elem_list[0].get_attribute('datetime')
            except Exception:
                pass
            else:
                # print('considering post with timestamp_str: ', timestamp_str)
                if timestamp_str not in self._read_posts_timestamp_list:
                    self._unread_posts_element_list.append(elem)
                    # print('attached post to list')
                else:
                    # print('post already in list')
                    pass

    def unread_posts_available(self) -> bool:
        if len(self._unread_posts_element_list) == 0:
            self._update_unread_posts()
        return len(self._unread_posts_element_list) > 0
    
    def number_of_posts_read(self) -> int:
        return len(self._read_posts_timestamp_list)
    
    def _handle_popup(self):
        elem_list = self.webdriver.find_elements(By.XPATH, "//*[text()='Not now']")
        if len(elem_list) > 0:
            elem_list[0].click()

    def _accept_cookies(self):
        self._handle_popup()
        try:
            elem = self.webdriver.find_element(By.XPATH, "//*[text()='Accept all cookies']")
            elem.click()
        except Exception:
            logging.warning('failed accepting cookies')
            
    def _open_page(self, account: str, date_begin: datetime, date_end: datetime) -> None:
        url = self._url(account,date_begin, date_end)
        self.webdriver.get(str(url))
        # wait for popup to appear and click on 'not now'-button
        try:
            elem = WebDriverWait(
                driver=self.webdriver, 
                timeout=60
            ).until(
                EC.presence_of_element_located(
                    (By.XPATH, "//*[text()='Not now']")
                )
            )
        except TimeoutError:
            pass
        else:
            elem.click()
        # click allow cookies
        self._accept_cookies()
        # identify timeline element
        self.timeline_elem = self.webdriver.find_element(By.XPATH, """//div[@aria-label="Timeline: Search timeline"]""")
        # update unread posts
        self._update_unread_posts()

    def _extract_account(self, elem):
        self._handle_popup()
        try:
            elem_list = elem.find_elements(By.XPATH, ".//span[contains(text(), '@')]")
            assert len(elem_list) > 0, 'failed locating account'
            self.actiondriver.move_to_element(elem_list[0]).perform()
            account_str = elem_list[0].text
        except Exception:
            logging.error('failed extracting account name')
            account_str = ''
        return account_str

    def _extract_time(self, elem):
        self._handle_popup()
        try:
            elem_list = elem.find_elements(By.XPATH, './/time')
            assert len(elem_list) > 0, 'failed locating time of post'
            self.actiondriver.move_to_element(elem_list[0]).perform()
            datetime_str = elem_list[0].get_attribute('datetime')
        except Exception:
            logging.error('failed extracting post time')
            datetime_str = ''
        return datetime_str

    def _extract_text(self, elem):
        self._handle_popup()
        try:
            elem_list = elem.find_elements(By.XPATH, './/div[@data-testid="tweetText"]')
            assert len(elem_list) > 0, 'failed locating text of post'
            self.actiondriver.move_to_element(elem_list[0]).perform()
            text = elem_list[0].text
        except Exception:
            logging.error('failed extracting post text')
            text = ''
        return text
    
    def _extract_post_url(self, elem):
        self._handle_popup()
        try:
            share_button = elem.find_element(By.XPATH, """.//div[@aria-label="Share Tweet"]""")
            # WebDriverWait(self.webdriver, 15).until(EC.element_to_be_clickable((By.XPATH, "//div[@aria-label='Share Tweet']"))).click()
            # self.webdriver.execute_script("arguments[0].scrollIntoView();", elem_list[0])
            # scroll to element
            self.webdriver.execute_script(
                """arguments[0].scrollIntoView({behavior: "smooth", block: "center", inline: "nearest"})""", 
                share_button
            )
            # move mouse to element
            self.actiondriver.move_to_element(share_button).perform()
            # click element
            share_button = elem.find_element(By.XPATH, """.//div[@aria-label="Share Tweet"]""")
            share_button.click()
            # wait for popup button to appear
            share_button = WebDriverWait(
                driver=self.webdriver, 
                timeout=10
            ).until(
                EC.presence_of_element_located(
                    (By.XPATH, """//span[contains(text(), "Copy link to Tweet")]""")
                )
            )
            # click button
            # elem_list = elem.find_elements(By.XPATH, "//span[contains(text(), 'Copy link to Tweet')]")
            self.actiondriver.move_to_element(share_button).perform()
            share_button.click()
            # paste copied url
            url = pyperclip.paste()
        except Exception:
            logging.error('failed extracting post url')
            url = ''
        return url
    
    def _extract_video_url(self, elem):
        self._handle_popup()
        try:
            elem_list = elem.find_elements(By.TAG_NAME, 'video')
            if len(elem_list) > 0:
                video_url = elem_list[0].get_attribute('src')
            else:
                video_url = ''
        except Exception:
            logging.error('failed extracting post video')
            video_url = ''
        return video_url
    
    def _extract_post_images(self, elem):
        self._handle_popup()
        img_list = list()
        try:
            elem_list = elem.find_elements(By.TAG_NAME, 'img')
            if len(elem_list) == 0:
                return img_list
            # get image urls
            image_url_list = [e.get_attribute('src') for e in elem_list]
            # keep only media-related urls
            image_url_list = [url for url in image_url_list if 'media' in url]
            if len(image_url_list) == 0:
                return img_list
            # download image data
            for url in image_url_list:
                try:
                    data = requests.get(url).content
                except Exception:
                    logging.error(f'failed downloading image: {url}')
                    continue
                img_data = BytesIO(data)
                img_list.append(img_data)
        except Exception:
            logging.error('failed extracting post images')
        return img_list

    def get_post(self):
        if len(self._unread_posts_element_list) == 0:
            raise ValueError('no unread posts to read')
        # get element
        elem = self._unread_posts_element_list.pop(0)
        # prepare response
        content = dict()
        try:
            # look for popup
            self._handle_popup()
            # scroll to element
            self.actiondriver.move_to_element(elem).perform()
            # get account
            content['account'] = self._extract_account(elem)
            # get time
            datetime_str = self._extract_time(elem)
            # print(datetime_str)
            content['time'] = datetime.strptime(datetime_str, '%Y-%m-%dT%H:%M:%S.%fZ')
            # get text
            content['text'] = self._extract_text(elem)
            # get post link
            content['url'] = self._extract_post_url(elem)
            # get video url
            content['video'] = self._extract_video_url(elem)
            # get image(s)
            content['images'] = self._extract_post_images(elem)
        except KeyboardInterrupt:
            return
        except Exception:
            traceback.print_exc()
        else:
            # add timestamp_str to list of read posts
            self._read_posts_timestamp_list.append(datetime_str)
            # print('sucessfully read post')
        return content, elem

def generate_date_list(until_year=2010):
    end_date = datetime.utcnow().replace(hour=0, minute=0, second=0, microsecond=0)
    start_date = end_date - timedelta(days=10)
    end_date_list = [end_date]
    start_date_list = [start_date]
    while start_date.year > until_year:
        end_date = start_date - timedelta(days=1)
        end_date_list.append(end_date)
        start_date = end_date - timedelta(days=10)
        start_date_list.append(start_date)
    date_list = list(zip(start_date_list, end_date_list))
    return date_list

def get_period(account: str, start_date: datetime, end_date: datetime):
    try:
        # setup
        twitter = Twitter()
        twitter.setup()
        twitter._open_page(account, start_date, end_date)
        db = connect_mongo()
        # scrape
        limit = 500
        attempts = 0
        while attempts < limit and twitter.unread_posts_available():
            content, elem = twitter.get_post()
            # print(content)
            t = Tweet(**content)
            db_id = db.insert_one(t.as_db_dict()).inserted_id
            # print(f"""inserted tweet from {t.time.strftime('%Y-%m-%d %H:%M:%S')} with db id: {db_id}""")
            attempts += 1
        num_posts_read = twitter.number_of_posts_read()
        # print('number of posts read: ', num_posts_read)
    except Exception as e:
        raise e
    finally:
        twitter.close()
    return num_posts_read

def scrape_account(account: str, until_year: int):
    # generate date list
    date_list = generate_date_list(until_year)
    # start scraping
    i = 0
    limit = 1000
    while len(date_list) > 0:
        if i >= limit:
            logging.error('scraping attempt limit reached!')
            return
        # get period
        start_date, end_date = date_list.pop(0)
        try:
            num_posts_found = get_period(account, start_date, end_date)
        except KeyboardInterrupt:
            logging.warning('KeyboardInterrupt')
            return
        except:
            date_list.append((start_date, end_date))
            logging.warning(f'retrying later due to error getting period {start_date} to {end_date}')
        else:
            logging.info(f'found {num_posts_found} posts between {start_date} and {end_date}')
        i += 1

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    account = 'Maersk'
    until_year = 2010
    scrape_account(account, until_year)