add type hints

Author: Yicong-Huang

pyspark_datasources/robinhood.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Dict
+from typing import Dict, List, Optional, Generator, Union
 import requests
 import json
 import base64
@@ -10,67 +10,9 @@
 from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition
 
 
-class RobinhoodDataSource(DataSource):
-    """
-    A data source for reading cryptocurrency data from Robinhood Crypto API.
-
-    This data source allows you to fetch real-time cryptocurrency market data,
-    trading pairs, and price information using Robinhood's official Crypto API.
-    It implements proper API key authentication and signature-based security.
-
-    Name: `robinhood`
-
-    Schema: `symbol string, price double, bid_price double, ask_price double, updated_at string`
-
-    Examples
-    --------
-    Register the data source:
-
-    >>> from pyspark_datasources import RobinhoodDataSource
-    >>> spark.dataSource.register(RobinhoodDataSource)
-
-    Load cryptocurrency market data with API authentication:
-
-    >>> df = spark.read.format("robinhood") \\
-    ...     .option("api_key", "your-api-key") \\
-    ...     .option("private_key", "your-base64-private-key") \\
-    ...     .load("BTC-USD,ETH-USD,DOGE-USD")
-    >>> df.show()
-    +--------+--------+---------+---------+--------------------+
-    |  symbol|   price|bid_price|ask_price|          updated_at|
-    +--------+--------+---------+---------+--------------------+
-    |BTC-USD |45000.50|45000.25 |45000.75 |2024-01-15T16:00:...|
-    |ETH-USD | 2650.75| 2650.50 | 2651.00 |2024-01-15T16:00:...|
-    |DOGE-USD|   0.085|    0.084|    0.086|2024-01-15T16:00:...|
-    +--------+--------+---------+---------+--------------------+
 
 
 
-    Notes
-    -----
-    - Requires valid Robinhood Crypto API credentials (API key and base64-encoded private key)
-    - Supports all major cryptocurrencies available on Robinhood
-    - Implements proper API authentication with NaCl (Sodium) signing
-    - Rate limiting is handled automatically
-    - Based on official Robinhood Crypto Trading API documentation
-    - Requires 'pynacl' library for cryptographic signing: pip install pynacl
-    - Reference: https://docs.robinhood.com/crypto/trading/
-    """
-
-    @classmethod
-    def name(cls) -> str:
-        return "robinhood"
-
-    def schema(self) -> str:
-        return (
-            "symbol string, price double, bid_price double, ask_price double, "
-            "updated_at string"
-        )
-
-    def reader(self, schema: StructType):
-        return RobinhoodDataReader(schema, self.options)
-
-
 @dataclass
 class CryptoPair(InputPartition):
     """Represents a single crypto trading pair partition for parallel processing."""
@@ -80,7 +22,7 @@ class CryptoPair(InputPartition):
 class RobinhoodDataReader(DataSourceReader):
     """Reader implementation for Robinhood Crypto API data source."""
 
-    def __init__(self, schema: StructType, options: Dict):
+    def __init__(self, schema: StructType, options: Dict[str, str]) -> None:
         self.schema = schema
         self.options = options
 
@@ -131,7 +73,7 @@ def _generate_signature(self, timestamp: int, method: str, path: str, body: str
         signature = base64.b64encode(signed.signature).decode("utf-8")
         return signature
 
-    def _make_authenticated_request(self, method: str, path: str, params: Dict = None, json_data: Dict = None):
+    def _make_authenticated_request(self, method: str, path: str, params: Optional[Dict[str, str]] = None, json_data: Optional[Dict] = None) -> Optional[Dict]:
         """Make an authenticated request to the Robinhood Crypto API."""
         timestamp = self._get_current_timestamp()
         url = self.base_url + path
@@ -168,14 +110,14 @@ def _make_authenticated_request(self, method: str, path: str, params: Dict = Non
             return None
 
     @staticmethod
-    def _get_query_params(key: str, *args) -> str:
+    def _get_query_params(key: str, *args: str) -> str:
         """Build query parameters for API requests."""
         if not args:
             return ""
         params = [f"{key}={arg}" for arg in args if arg]
         return "?" + "&".join(params)
 
-    def partitions(self):
+    def partitions(self) -> List[CryptoPair]:
         """Create partitions for parallel processing of crypto pairs."""
         # Use specified symbols from path
         symbols_str = self.options.get("path", "")
@@ -196,7 +138,7 @@ def partitions(self):
 
         return [CryptoPair(symbol=symbol) for symbol in formatted_symbols]
 
-    def read(self, partition: CryptoPair):
+    def read(self, partition: CryptoPair) -> Generator[Row, None, None]:
         """Read crypto data for a single trading pair partition."""
         symbol = partition.symbol
 
@@ -206,7 +148,7 @@ def read(self, partition: CryptoPair):
             # Log error but don't fail the entire job
             print(f"Warning: Failed to fetch data for {symbol}: {str(e)}")
 
-    def _read_crypto_pair_data(self, symbol: str):
+    def _read_crypto_pair_data(self, symbol: str) -> Generator[Row, None, None]:
         """Fetch cryptocurrency market data for a given trading pair."""
         try:
             # Get best bid/ask data for the trading pair using query parameters
@@ -216,7 +158,7 @@ def _read_crypto_pair_data(self, symbol: str):
             if market_data and 'results' in market_data:
                 for quote in market_data['results']:
                     # Parse numeric values safely
-                    def safe_float(value, default=0.0):
+                    def safe_float(value: Union[str, int, float, None], default: float = 0.0) -> float:
                         if value is None or value == "":
                             return default
                         try:
@@ -246,3 +188,64 @@ def safe_float(value, default=0.0):
             print(f"Data parsing error for {symbol}: {str(e)}")
         except Exception as e:
             print(f"Unexpected error fetching data for {symbol}: {str(e)}")
+
+
+class RobinhoodDataSource(DataSource):
+    """
+    A data source for reading cryptocurrency data from Robinhood Crypto API.
+
+    This data source allows you to fetch real-time cryptocurrency market data,
+    trading pairs, and price information using Robinhood's official Crypto API.
+    It implements proper API key authentication and signature-based security.
+
+    Name: `robinhood`
+
+    Schema: `symbol string, price double, bid_price double, ask_price double, updated_at string`
+
+    Examples
+    --------
+    Register the data source:
+
+    >>> from pyspark_datasources import RobinhoodDataSource
+    >>> spark.dataSource.register(RobinhoodDataSource)
+
+    Load cryptocurrency market data with API authentication:
+
+    >>> df = spark.read.format("robinhood") \\
+    ...     .option("api_key", "your-api-key") \\
+    ...     .option("private_key", "your-base64-private-key") \\
+    ...     .load("BTC-USD,ETH-USD,DOGE-USD")
+    >>> df.show()
+    +--------+--------+---------+---------+--------------------+
+    |  symbol|   price|bid_price|ask_price|          updated_at|
+    +--------+--------+---------+---------+--------------------+
+    |BTC-USD |45000.50|45000.25 |45000.75 |2024-01-15T16:00:...|
+    |ETH-USD | 2650.75| 2650.50 | 2651.00 |2024-01-15T16:00:...|
+    |DOGE-USD|   0.085|    0.084|    0.086|2024-01-15T16:00:...|
+    +--------+--------+---------+---------+--------------------+
+
+
+
+    Notes
+    -----
+    - Requires valid Robinhood Crypto API credentials (API key and base64-encoded private key)
+    - Supports all major cryptocurrencies available on Robinhood
+    - Implements proper API authentication with NaCl (Sodium) signing
+    - Rate limiting is handled automatically
+    - Based on official Robinhood Crypto Trading API documentation
+    - Requires 'pynacl' library for cryptographic signing: pip install pynacl
+    - Reference: https://docs.robinhood.com/crypto/trading/
+    """
+
+    @classmethod
+    def name(cls) -> str:
+        return "robinhood"
+
+    def schema(self) -> str:
+        return (
+            "symbol string, price double, bid_price double, ask_price double, "
+            "updated_at string"
+        )
+
+    def reader(self, schema: StructType) -> RobinhoodDataReader:
+        return RobinhoodDataReader(schema, self.options)