refactor: improve typing

Author: ali-behjati

README.md

@@ -14,7 +14,7 @@ Install dependencies and run the service:
 ```sh
 $ poetry env use $(which python) # point Poetry to the pyenv python shim
 $ poetry install
-$ poetry run pyth-observer
+$ poetry run pyth-observe --config sample.config.yaml --publishers sample.publishers.yaml --coingecko-mapping sample.coingecko.yaml
 ```
 
 Use `poetry run pyth-observer --help` for documentation on arguments and environment variables.

pyth_observer/__init__.py

@@ -1,6 +1,6 @@
 import asyncio
 import os
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Literal, Tuple
 
 from base58 import b58decode
 from loguru import logger
@@ -18,6 +18,7 @@
 from throttler import Throttler
 
 import pyth_observer.health_server as health_server
+from pyth_observer.check import State
 from pyth_observer.check.price_feed import PriceFeedState
 from pyth_observer.check.publisher import PublisherState
 from pyth_observer.coingecko import Symbol, get_coingecko_prices
@@ -35,7 +36,9 @@
 ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
 
 
-def get_solana_urls(network) -> Tuple[str, str]:
+def get_solana_urls(
+    network: Literal["devnet", "testnet", "mainnet", "pythtest", "pythnet"]
+) -> Tuple[str, str]:
     """
     Helper for getting the correct urls for the PythClient
     """
@@ -55,7 +58,7 @@ def __init__(
         config: Dict[str, Any],
         publishers: Dict[str, Publisher],
         coingecko_mapping: Dict[str, Symbol],
-    ):
+    ) -> None:
         self.config = config
         self.dispatch = Dispatch(config, publishers)
         self.publishers = publishers
@@ -77,9 +80,9 @@ def __init__(
             config=config,
         )
 
-    async def run(self):
+    async def run(self) -> None:
         # global states
-        states = []
+        states: List[State] = []
         while True:
             try:
                 logger.info("Running checks")
@@ -91,7 +94,7 @@ async def run(self):
                 health_server.observer_ready = True
 
                 processed_feeds = 0
-                active_publishers_by_symbol = {}
+                active_publishers_by_symbol: Dict[str, Dict[str, Any]] = {}
 
                 for product in products:
                     # Skip tombstone accounts with blank metadata
@@ -104,7 +107,7 @@ async def run(self):
                     # For each product, we build a list of price feed states (one
                     # for each price account) and a list of publisher states (one
                     # for each publisher).
-                    states = []
+                    states: List[State] = []
                     price_accounts = await self.get_pyth_prices(product)
 
                     crosschain_price = crosschain_prices.get(
@@ -249,7 +252,9 @@ async def get_pyth_prices(
             ).inc()
             raise
 
-    async def get_coingecko_prices(self):
+    async def get_coingecko_prices(
+        self,
+    ) -> Tuple[Dict[str, float], Dict[str, int]]:
         logger.debug("Fetching CoinGecko prices...")
 
         try:

pyth_observer/check/price_feed.py

@@ -1,7 +1,7 @@
 import time
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Dict, Optional, Protocol, runtime_checkable
+from typing import Any, Dict, Optional, Protocol, runtime_checkable
 from zoneinfo import ZoneInfo
 
 import arrow
@@ -33,7 +33,7 @@ class PriceFeedState:
 
 @runtime_checkable
 class PriceFeedCheck(Protocol):
-    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig):
+    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig) -> None:
         ...
 
     def state(self) -> PriceFeedState:
@@ -42,12 +42,12 @@ def state(self) -> PriceFeedState:
     def run(self) -> bool:
         ...
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         ...
 
 
 class PriceFeedOfflineCheck(PriceFeedCheck):
-    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig):
+    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig) -> None:
         self.__state = state
         self.__max_slot_distance: int = int(config["max_slot_distance"])
         self.__abandoned_slot_distance: int = int(config["abandoned_slot_distance"])
@@ -79,7 +79,7 @@ def run(self) -> bool:
         # Fail
         return False
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         distance = self.__state.latest_block_slot - self.__state.latest_trading_slot
         return {
             "msg": f"{self.__state.symbol} is offline (either non-trading/stale). Last update {distance} slots ago.",
@@ -91,7 +91,7 @@ def error_message(self) -> dict:
 
 
 class PriceFeedCoinGeckoCheck(PriceFeedCheck):
-    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig):
+    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig) -> None:
         self.__state = state
         self.__max_deviation: int = int(config["max_deviation"])  # Percentage
         self.__max_staleness: int = int(config["max_staleness"])  # Seconds
@@ -124,7 +124,7 @@ def run(self) -> bool:
         # Fail
         return False
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         return {
             "msg": f"{self.__state.symbol} is too far from Coingecko's price.",
             "type": "PriceFeedCoinGeckoCheck",
@@ -135,7 +135,7 @@ def error_message(self) -> dict:
 
 
 class PriceFeedConfidenceIntervalCheck(PriceFeedCheck):
-    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig):
+    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig) -> None:
         self.__state = state
         self.__min_confidence_interval: int = int(config["min_confidence_interval"])
 
@@ -154,7 +154,7 @@ def run(self) -> bool:
         # Fail
         return False
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         return {
             "msg": f"{self.__state.symbol} confidence interval is too low.",
             "type": "PriceFeedConfidenceIntervalCheck",
@@ -164,7 +164,7 @@ def error_message(self) -> dict:
 
 
 class PriceFeedCrossChainOnlineCheck(PriceFeedCheck):
-    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig):
+    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig) -> None:
         self.__state = state
         self.__max_staleness: int = int(config["max_staleness"])
 
@@ -204,7 +204,7 @@ def run(self) -> bool:
         # Fail
         return False
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         if self.__state.crosschain_price:
             publish_time = arrow.get(self.__state.crosschain_price["publish_time"])
         else:
@@ -219,7 +219,7 @@ def error_message(self) -> dict:
 
 
 class PriceFeedCrossChainDeviationCheck(PriceFeedCheck):
-    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig):
+    def __init__(self, state: PriceFeedState, config: PriceFeedCheckConfig) -> None:
         self.__state = state
         self.__max_deviation: int = int(config["max_deviation"])
         self.__max_staleness: int = int(config["max_staleness"])
@@ -262,7 +262,7 @@ def run(self) -> bool:
         # Fail
         return False
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         # It can never happen because of the check logic but linter could not understand it.
         price = (
             self.__state.crosschain_price["price"]

pyth_observer/check/publisher.py

@@ -2,7 +2,7 @@
 from collections import defaultdict, deque
 from dataclasses import asdict, dataclass
 from datetime import datetime
-from typing import Dict, Protocol, runtime_checkable
+from typing import Any, Dict, Protocol, runtime_checkable
 from zoneinfo import ZoneInfo
 
 from loguru import logger
@@ -54,7 +54,7 @@ class PublisherState:
 
 @runtime_checkable
 class PublisherCheck(Protocol):
-    def __init__(self, state: PublisherState, config: PublisherCheckConfig):
+    def __init__(self, state: PublisherState, config: PublisherCheckConfig) -> None:
         ...
 
     def state(self) -> PublisherState:
@@ -63,12 +63,12 @@ def state(self) -> PublisherState:
     def run(self) -> bool:
         ...
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         ...
 
 
 class PublisherWithinAggregateConfidenceCheck(PublisherCheck):
-    def __init__(self, state: PublisherState, config: PublisherCheckConfig):
+    def __init__(self, state: PublisherState, config: PublisherCheckConfig) -> None:
         self.__state = state
         self.__max_interval_distance: int = int(config["max_interval_distance"])
 
@@ -103,7 +103,7 @@ def run(self) -> bool:
         # Fail
         return False
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         diff = self.__state.price - self.__state.price_aggregate
         intervals_away = abs(diff / self.__state.confidence_interval_aggregate)
         return {
@@ -117,7 +117,7 @@ def error_message(self) -> dict:
 
 
 class PublisherConfidenceIntervalCheck(PublisherCheck):
-    def __init__(self, state: PublisherState, config: PublisherCheckConfig):
+    def __init__(self, state: PublisherState, config: PublisherCheckConfig) -> None:
         self.__state = state
         self.__min_confidence_interval: int = int(config["min_confidence_interval"])
 
@@ -141,7 +141,7 @@ def run(self) -> bool:
         # Fail
         return False
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         return {
             "msg": f"{self.__state.publisher_name} confidence interval is too tight.",
             "type": "PublisherConfidenceIntervalCheck",
@@ -153,7 +153,7 @@ def error_message(self) -> dict:
 
 
 class PublisherOfflineCheck(PublisherCheck):
-    def __init__(self, state: PublisherState, config: PublisherCheckConfig):
+    def __init__(self, state: PublisherState, config: PublisherCheckConfig) -> None:
         self.__state = state
         self.__max_slot_distance: int = int(config["max_slot_distance"])
         self.__abandoned_slot_distance: int = int(config["abandoned_slot_distance"])
@@ -182,7 +182,7 @@ def run(self) -> bool:
         # Fail
         return False
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         distance = self.__state.latest_block_slot - self.__state.slot
         return {
             "msg": f"{self.__state.publisher_name} hasn't published recently for {distance} slots.",
@@ -195,7 +195,7 @@ def error_message(self) -> dict:
 
 
 class PublisherPriceCheck(PublisherCheck):
-    def __init__(self, state: PublisherState, config: PublisherCheckConfig):
+    def __init__(self, state: PublisherState, config: PublisherCheckConfig) -> None:
         self.__state = state
         self.__max_aggregate_distance: int = int(config["max_aggregate_distance"])  # %
         self.__max_slot_distance: int = int(config["max_slot_distance"])  # Slots
@@ -230,7 +230,7 @@ def run(self) -> bool:
         # Fail
         return False
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         deviation = (self.ci_adjusted_price_diff() / self.__state.price_aggregate) * 100
         return {
             "msg": f"{self.__state.publisher_name} price is too far from aggregate price.",
@@ -250,7 +250,7 @@ def ci_adjusted_price_diff(self) -> float:
 
 
 class PublisherStalledCheck(PublisherCheck):
-    def __init__(self, state: PublisherState, config: PublisherCheckConfig):
+    def __init__(self, state: PublisherState, config: PublisherCheckConfig) -> None:
         self.__state = state
         self.__stall_time_limit: int = int(
             config["stall_time_limit"]
@@ -313,7 +313,7 @@ def run(self) -> bool:
 
         return not result.is_stalled
 
-    def error_message(self) -> dict:
+    def error_message(self) -> Dict[str, Any]:
         stall_duration = f"{self.__last_analysis.duration:.1f} seconds"
         return {
             "msg": f"{self.__state.publisher_name} has been publishing the same price of {self.__state.symbol} for {stall_duration}",

pyth_observer/check/stall_detection.py

@@ -64,7 +64,7 @@ def __init__(
         stall_time_limit: float,
         noise_threshold: float = 1e-4,
         min_noise_samples: int = 5,
-    ):
+    ) -> None:
         """
         Initialize stall detector.
 

pyth_observer/cli.py

@@ -1,6 +1,7 @@
 import asyncio
 import os
 import sys
+from typing import Any, Dict
 
 import click
 import yaml
@@ -37,11 +38,13 @@
     envvar="PROMETHEUS_PORT",
     default="9001",
 )
-def run(config, publishers, coingecko_mapping, prometheus_port):
-    config_ = yaml.safe_load(open(config, "r"))
+def run(
+    config: str, publishers: str, coingecko_mapping: str, prometheus_port: str
+) -> None:
+    config_: Dict[str, Any] = yaml.safe_load(open(config, "r"))  # type: ignore[assignment]
     # Load publishers YAML file and convert to dictionary of Publisher instances
-    publishers_raw = yaml.safe_load(open(publishers, "r"))
-    publishers_ = {
+    publishers_raw: list[Dict[str, Any]] = yaml.safe_load(open(publishers, "r"))  # type: ignore[assignment]
+    publishers_: Dict[str, Publisher] = {
         publisher["key"]: Publisher(
             key=publisher["key"],
             name=publisher["name"],
@@ -53,7 +56,7 @@ def run(config, publishers, coingecko_mapping, prometheus_port):
         )
         for publisher in publishers_raw
     }
-    coingecko_mapping_ = yaml.safe_load(open(coingecko_mapping, "r"))
+    coingecko_mapping_: Dict[str, Any] = yaml.safe_load(open(coingecko_mapping, "r"))  # type: ignore[assignment]
     observer = Observer(
         config_,
         publishers_,
@@ -62,7 +65,7 @@ def run(config, publishers, coingecko_mapping, prometheus_port):
 
     start_http_server(int(prometheus_port))
 
-    async def main():
+    async def main() -> None:
         asyncio.create_task(start_health_server())
         await observer.run()
 

pyth_observer/coingecko.py

@@ -1,4 +1,4 @@
-from typing import Dict, TypedDict
+from typing import Any, Dict, TypedDict
 
 from loguru import logger
 from pycoingecko import CoinGeckoAPI
@@ -15,7 +15,9 @@ class Symbol(TypedDict):
 # However prices are updated every 1-10 minutes: https://www.coingecko.com/en/faq
 # Hence we only have to query once every minute.
 @throttle(rate_limit=1, period=60)
-async def get_coingecko_prices(mapping: Dict[str, Symbol]):
+async def get_coingecko_prices(
+    mapping: Dict[str, Symbol],
+) -> Dict[str, Dict[str, Any]]:
     inverted_mapping = {mapping[x]["api"]: x for x in mapping}
     ids = [mapping[x]["api"] for x in mapping]
 

pyth_observer/crosschain.py

@@ -16,7 +16,7 @@ class CrosschainPrice(TypedDict):
 
 
 class CrosschainPriceObserver:
-    def __init__(self, url):
+    def __init__(self, url: str) -> None:
         self.url = url
         self.valid = self.is_endpoint_valid()