improve fetchers

Author: danielmoessner

apps/stocks/fetcher/__init__.py

@@ -1,11 +1,9 @@
 from typing import Callable
 from apps.stocks.fetcher.marketstack import (
-    fetch_price_with_marketstack_fetcher,
-    fetch_prices_with_marketstack_fetchers,
+    MarketstackFetcher,
 )
 from apps.stocks.fetcher.selenium import (
-    fetch_price_with_selenium_fetcher,
-    fetch_prices_with_selenium_fetchers,
+    SeleniumFetcher,
 )
 from apps.stocks.fetcher.website import (
     WebsiteFetcher,
@@ -16,6 +14,22 @@
 from typing import Callable
 from django.utils import timezone
 from apps.stocks.models import Price, PriceFetcher, Stock
+from django.utils import timezone
+from apps.stocks.models import Stock
+from apps.stocks.forms import PriceForm
+
+
+def save_price(price: float, stock: Stock) -> None:
+    price = PriceForm(
+        {
+            "ticker": stock.ticker,
+            "exchange": stock.exchange,
+            "date": timezone.now(),
+            "price": price,
+        }
+    )
+    if price.is_valid():
+        price.save()
 
 
 FETCHER_FUNCTION = Callable[[PriceFetcher], tuple[bool, str]]
@@ -32,25 +46,47 @@ def get_stocks_to_be_fetched() -> list[Stock]:
     return stocks_to_be_fetched
 
 
-def get_fetchers_to_be_run(type: str) -> list[PriceFetcher]:
+def get_fetchers_to_be_run(type: str) -> dict[str, dict[str, int | str]]:
     fetchers_to_be_run = []
     for fetcher in list(PriceFetcher.objects.filter(type=type)):
         if Price.objects.filter(
             ticker=fetcher.stock.ticker, date__gt=timezone.now() - timedelta(days=1)
         ).exists():
             continue
         fetchers_to_be_run.append(fetcher)
-    return fetchers_to_be_run
+    return {str(fetcher.pk): fetcher.data for fetcher in fetchers_to_be_run}
+
+
+def turn_fetchers_to_data(fetchers: list[PriceFetcher]) -> dict[str, dict]:
+    data = {}
+    for fetcher in fetchers:
+        data[str(fetcher.pk)] = fetcher.data
+    return data
+
+
+def save_prices(results: dict[str, tuple[bool, str | float]]):
+    for fetcher, result in results:
+        fetcher = PriceFetcher.objects.get(pk=fetcher)
+        if result[0]:
+            save_price(fetcher.stock, result[1])
 
 
 FETCHERS: dict[str, FETCHER_FUNCTION] = {
     "WEBSITE": WebsiteFetcher.fetch_single,
-    "SELENIUM": fetch_price_with_selenium_fetcher,
-    "MARKETSTACK": fetch_price_with_marketstack_fetcher,
+    "SELENIUM": SeleniumFetcher.fetch_single,
+    "MARKETSTACK": MarketstackFetcher.fetch_single,
 }
 
 
 def fetch_prices():
-    WebsiteFetcher.fetch_multiple(get_fetchers_to_be_run("WEBSITE"))
-    fetch_prices_with_selenium_fetchers(get_stocks_to_be_fetched())
-    fetch_prices_with_marketstack_fetchers(get_stocks_to_be_fetched())
+    fetchers = get_fetchers_to_be_run("WEBSITE")
+    results = WebsiteFetcher.fetch_multiple(fetchers)
+    save_prices(results)
+
+    fetchers = get_fetchers_to_be_run("SELENIUM")
+    results = SeleniumFetcher.fetch_multiple(fetchers)
+    save_prices(results)
+
+    fetchers = get_fetchers_to_be_run("MARKETSTACK")
+    results = MarketstackFetcher.fetch_multiple(fetchers)
+    save_prices(results)

apps/stocks/fetcher/base.py

@@ -1,13 +1,13 @@
 import abc
 
-from apps.stocks.models import PriceFetcher, Stock
-
 
 class Fetcher(abc.ABC):
     @abc.abstractmethod
-    def fetch_single(self, stock: PriceFetcher) -> tuple[bool, str | float]:
+    def fetch_single(self, *args, **kwargs) -> tuple[bool, str | float]:
         raise NotImplementedError()
-    
+
     @abc.abstractmethod
-    def fetch_multiple(self, stocks: list[PriceFetcher]) -> dict[Stock, tuple[bool, str | float]]:
+    def fetch_multiple(
+        self, *args, **kwargs
+    ) -> dict[str, tuple[bool, str | float]]:
         raise NotImplementedError()

apps/stocks/fetcher/marketstack.py

@@ -1,43 +1,39 @@
 from django.conf import settings
 import requests
-from apps.stocks.fetcher.price import save_price
-from apps.stocks.models import PriceFetcher, Stock
-
-
-def fetch_prices_from_marketstack(symbols: list[str]) -> tuple[bool, str]:
-    symbols = ",".join(symbols)
-    params = {"access_key": settings.MARKETSTACK_API_KEY}
-    url = "http://api.marketstack.com/v1/eod/latest?symbols={}".format(symbols)
-    api_result = requests.get(url, params)
-    api_response = api_result.json()
-    if "error" in api_response:
-        return (
-            False,
-            f"Could not fetch prices from marketstack: '{api_response['error']['message']}'.",
-        )
-
-    for price in api_response["data"]:
-        save_price(
-            round(price["close"], 2), Stock.objects.get(ticker=[price["symbol"]])
-        )
-
-    return True, ""
-
-
-def fetch_price_with_marketstack_fetcher(fetcher: PriceFetcher) -> tuple[bool, str]:
-    return fetch_prices_from_marketstack([fetcher.data["symbol"]])
-
-
-def fetch_prices_with_marketstack_fetchers(
-    stocks: list[Stock], messages: list[str] | None = None
-):
-    symbols = []
-    for stock in stocks:
-        fetcher = stock.price_fetchers.filter(type="MARKETSTACK").first()
-        if fetcher is None:
-            if messages is not None:
-                messages.append(f"Could not find a price fetcher for {stock.ticker}.")
-            continue
-        symbols.append(fetcher.data["symbol"])
-
-    fetch_prices_from_marketstack(symbols)
+from apps.stocks.fetcher.base import Fetcher
+
+
+class MarketstackFetcher(Fetcher):
+    def fetch_single(self, symbol: str) -> tuple[bool, str | float]:
+        return self.fetch_multiple({"_": {"symbol": symbol}})["_"]
+
+    def fetch_multiple(
+        self, data: dict[str, dict[str, str | int]]
+    ) -> dict[str, tuple[bool, str | float]]:
+        symbols = []
+        for fetcher, input in data.items():
+            symbols.append(input["symbol"])
+
+        symbols = ",".join(symbols)
+        params = {"access_key": settings.MARKETSTACK_API_KEY}
+        url = "http://api.marketstack.com/v1/eod/latest?symbols={}".format(symbols)
+        api_result = requests.get(url, params)
+        api_response = api_result.json()
+
+        results = {}
+
+        if "error" in api_response:
+            for fetcher in data:
+                results[fetcher] = (
+                    False,
+                    f"Could not fetch prices from marketstack: '{api_response['error']['message']}'.",
+                )
+            return results
+        
+        for fetcher, input in data.items():
+            symbol = input["symbol"]
+            results[fetcher] = (False, "Price not found in response.")
+            for price in api_response["data"]:
+                if price["symbol"] == symbol:
+                    results[fetcher] = (True, round(price["close"], 2))
+                    break

apps/stocks/fetcher/selenium.py

@@ -2,57 +2,43 @@
 import time
 from bs4 import BeautifulSoup
 from apps.core.selenium import get_chrome_driver
-from apps.stocks.fetcher.price import save_price
-from apps.stocks.models import PriceFetcher, Stock
-
-
-def fetch_price_with_selenium_fetcher(fetcher: PriceFetcher) -> tuple[bool, str]:
-    stock = fetcher.stock
-    website = fetcher.data["website"]
-    target = fetcher.data["target"]
-
-    browser = get_chrome_driver()
-    try:
-        browser.get(website)
-        time.sleep(5)  # wait for the api requests to finish
-        html = browser.page_source
-    except Exception as e:
-        return False, f"An error occured while trying to connect to {website}: {e}."
-    finally:
-        browser.quit()
-
-    soup = BeautifulSoup(html, features="html.parser")
-    selection = soup.select_one(target)
-
-    if not selection:
-        return (
-            False,
-            f"Could not find a price for {stock.ticker} on {website} with {target}.",
-        )
-
-    result = re.search("\d{1,5}[.,]\d{2}", str(selection))
-
-    if not result:
-        return False, f"Could not find a price inside '{selection}'."
-
-    price = result.group()
-    price = price.replace(",", ".")
-    price = float(price)
-    save_price(price, stock)
-    return True, ""
-
-
-def fetch_prices_with_selenium_fetchers(
-    stocks: list[Stock], messages: list[str] | None = None
-):
-    i = 0
-    for stock in stocks:
-        fetchers = list(stock.price_fetchers.filter(type="SELENIUM"))
-        for fetcher in fetchers:
-            success, message = fetch_price_with_selenium_fetcher(fetcher)
-            if success:
-                break
-            if messages is not None:
-                messages.append(message)
-        time.sleep(i)
-        i += 1
+from apps.stocks.fetcher.base import Fetcher
+
+
+class SeleniumFetcher(Fetcher):
+    def fetch_single(self, website: str, target: str) -> tuple[bool, str | float]:
+        browser = get_chrome_driver()
+        try:
+            browser.get(website)
+            time.sleep(5)  # wait for the api requests to finish
+            html = browser.page_source
+        except Exception as e:
+            return False, f"An error occured while trying to connect to {website}: {e}."
+        finally:
+            browser.quit()
+
+        soup = BeautifulSoup(html, features="html.parser")
+        selection = soup.select_one(target)
+
+        if not selection:
+            return (
+                False,
+                f"Could not find a price for on {website} with {target}.",
+            )
+
+        result = re.search("\d{1,5}[.,]\d{2}", str(selection))
+
+        if not result:
+            return False, f"Could not find a price inside '{selection}'."
+
+        price = result.group()
+        price = price.replace(",", ".")
+        price = float(price)
+        return True, price
+
+    def fetch_multiple(self, data: dict[str, dict[str, str | int]]) -> dict[str, tuple[bool, str | float]]:
+        results = {}
+        for fetcher, input in data.items():
+            result = self.fetch_single(**input)
+            results[fetcher] = result
+        return results

apps/stocks/fetcher/website.py

@@ -3,7 +3,6 @@
 from bs4 import BeautifulSoup
 import requests
 from apps.stocks.fetcher.base import Fetcher
-from apps.stocks.models import PriceFetcher, Stock
 
 
 headers = {
@@ -22,11 +21,7 @@
 
 
 class WebsiteFetcher(Fetcher):
-    def fetch_single(self, fetcher: PriceFetcher) -> tuple[bool, str | float]:
-        stock = fetcher.stock
-        website = fetcher.data["website"]
-        target = fetcher.data["target"]
-
+    def fetch_single(self, website: str, target: str) -> tuple[bool, str | float]:
         try:
             resp = requests.get(website, headers=headers)
             html = resp.text
@@ -45,7 +40,7 @@ def fetch_single(self, fetcher: PriceFetcher) -> tuple[bool, str | float]:
         if not selection:
             return (
                 False,
-                f"Could not find a price for {stock.ticker} on {website} with {target}.",
+                f"Could not find a price on {website} with {target}.",
             )
 
         result = re.search("\d{1,5}[.,]\d{2}", str(selection))
@@ -59,14 +54,13 @@ def fetch_single(self, fetcher: PriceFetcher) -> tuple[bool, str | float]:
         return True, price
 
     def fetch_multiple(
-        self, fetchers: list[PriceFetcher]
-    ) -> dict[Stock, tuple[bool, str | float]]:
+        self, data: dict[str, dict[str, str | int]]
+    ) -> dict[str, tuple[bool, str | float]]:
         i = 0
         results = {}
-        for fetcher in fetchers:
-            assert fetcher.type == "WEBSITE"
-            result = self.fetch_single(fetcher)
-            results[fetcher.stock] = result
+        for fetcher, input in data.items():
+            result = self.fetch_single(**input)
+            results[fetcher] = result
             time.sleep(i)
             i += 1
         return results