fixed bugs and ran tests

code/selenium_scraper.py

@@ -8,7 +8,7 @@ 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
+from datetime import datetime, timedelta
 from PIL import Image
 from io import BytesIO
 import requests
@@ -17,6 +17,8 @@ import logging
 import time
 import traceback
 
+from database import connect_mongo
+from classes import Tweet
 
 class Twitter(object):
     def __init__(self):
@@ -256,41 +258,69 @@ class Twitter(object):
             # 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 extract_tweet(webelement: WebElement) -> dict:
-#     post_time = webelement.find_element(By.XPATH, '//time').get_attribute('datetime')
-#     post_text = webelement.find_element(By.CSS_SELECTOR, '[data-testid="tweetText"]').text
-#     # find images in webelement
-#     img_url_list = [elem.get_attribute('src') for elem in webelement.find_elements(By.TAG_NAME, 'img')]
-#     # keep all images with url containing 'media'
-#     img_url_list = [url for url in img_url_list if 'media' in url]
-#     img_data_list = [requests.get(img_url).content for img_url in img_url_list]
-#     tweet = Tweet(
-#         time=post_time,
-#         text=post_text,
-#         images=img_data_list
-#     )
-#     return tweet
+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')
+        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
 
-# def fetch_period(account: str, date_begin: datetime, date_end: datetime) -> List[Tweet]:
-#     # prepare url
-#     url = base_url(account, date_begin, date_end)
-#     # setup webdriver
-#     driver = webdriver.Chrome('../chromedriver_mac')
-#     driver.get(str(url))
-#     # wait for popup to appear
-#     elem = WebDriverWait(
-#         driver=driver, 
-#         timeout=20
-#     ).until(
-#         EC.presence_of_element_located(
-#             (By.XPATH, "//*[text()='Not now']")
-#         )
-#     )
-#     # turn off notifications - click not now button
-#     elem.click()
-#     # click allow cookies
-#     elem = driver.find_element(By.XPATH, "//*[text()='Accept all cookies']")
-#     elem.click()
-    # 
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    account = 'Maersk'
+    until_year = 2010
+    scrape_account(account, until_year)