fix format

Author: Yicong-Huang

pyspark_datasources/robinhood.py

@@ -10,12 +10,10 @@
 from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition
 
 
-
-
-
 @dataclass
 class CryptoPair(InputPartition):
     """Represents a single crypto trading pair partition for parallel processing."""
+
     symbol: str
 
 
@@ -25,21 +23,22 @@ class RobinhoodDataReader(DataSourceReader):
     def __init__(self, schema: StructType, options: Dict[str, str]) -> None:
         self.schema = schema
         self.options = options
-        
+
         # Required API authentication
         self.api_key = options.get("api_key")
         self.private_key_base64 = options.get("private_key")
-        
+
         if not self.api_key or not self.private_key_base64:
             raise ValueError(
                 "Robinhood Crypto API requires both 'api_key' and 'private_key' options. "
                 "The private_key should be base64-encoded. "
                 "Get your API credentials from https://docs.robinhood.com/crypto/trading/"
             )
-        
+
         # Initialize NaCl signing key
         try:
             from nacl.signing import SigningKey
+
             private_key_seed = base64.b64decode(self.private_key_base64)
             self.signing_key = SigningKey(private_key_seed)
         except ImportError:
@@ -49,17 +48,14 @@ def __init__(self, schema: StructType, options: Dict[str, str]) -> None:
             )
         except Exception as e:
             raise ValueError(f"Invalid private key format: {str(e)}")
-        
-
 
-        
         # Crypto API base URL (configurable for testing)
         self.base_url = options.get("base_url", "https://trading.robinhood.com")
 
     def _get_current_timestamp(self) -> int:
         """Get current UTC timestamp."""
         return int(datetime.datetime.now(tz=datetime.timezone.utc).timestamp())
-    
+
     def _generate_signature(self, timestamp: int, method: str, path: str, body: str = "") -> str:
         """Generate NaCl signature for API authentication following Robinhood's specification."""
         # Official Robinhood signature format: f"{api_key}{current_timestamp}{path}{method}{body}"
@@ -68,41 +64,49 @@ def _generate_signature(self, timestamp: int, method: str, path: str, body: str
             message_to_sign = f"{self.api_key}{timestamp}{path}{method.upper()}"
         else:
             message_to_sign = f"{self.api_key}{timestamp}{path}{method.upper()}{body}"
-            
+
         signed = self.signing_key.sign(message_to_sign.encode("utf-8"))
         signature = base64.b64encode(signed.signature).decode("utf-8")
         return signature
 
-    def _make_authenticated_request(self, method: str, path: str, params: Optional[Dict[str, str]] = None, json_data: Optional[Dict] = None) -> Optional[Dict]:
+    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
-        
+
         # Prepare request body for signature (only for non-GET requests)
         body = ""
         if method.upper() != "GET" and json_data:
-            body = json.dumps(json_data, separators=(',', ':'))  # Compact JSON format
-        
+            body = json.dumps(json_data, separators=(",", ":"))  # Compact JSON format
+
         # Generate signature
         signature = self._generate_signature(timestamp, method, path, body)
-        
+
         # Set authentication headers
         headers = {
-            'x-api-key': self.api_key,
-            'x-signature': signature,
-            'x-timestamp': str(timestamp)
+            "x-api-key": self.api_key,
+            "x-signature": signature,
+            "x-timestamp": str(timestamp),
         }
-        
+
         try:
             # Make request
             if method.upper() == "GET":
                 response = requests.get(url, headers=headers, params=params, timeout=10)
             elif method.upper() == "POST":
-                headers['Content-Type'] = 'application/json'
+                headers["Content-Type"] = "application/json"
                 response = requests.post(url, headers=headers, json=json_data, timeout=10)
             else:
-                response = requests.request(method, url, headers=headers, params=params, json=json_data, timeout=10)
-            
+                response = requests.request(
+                    method, url, headers=headers, params=params, json=json_data, timeout=10
+                )
+
             response.raise_for_status()
             return response.json()
         except requests.RequestException as e:
@@ -116,32 +120,30 @@ def _get_query_params(key: str, *args: str) -> str:
             return ""
         params = [f"{key}={arg}" for arg in args if arg]
         return "?" + "&".join(params)
-    
+
     def partitions(self) -> List[CryptoPair]:
         """Create partitions for parallel processing of crypto pairs."""
         # Use specified symbols from path
         symbols_str = self.options.get("path", "")
         if not symbols_str:
-            raise ValueError(
-                "Must specify crypto pairs to load using .load('BTC-USD,ETH-USD')"
-            )
-            
+            raise ValueError("Must specify crypto pairs to load using .load('BTC-USD,ETH-USD')")
+
         # Split symbols by comma and create partitions
         symbols = [symbol.strip().upper() for symbol in symbols_str.split(",")]
         # Ensure proper format (e.g., BTC-USD)
         formatted_symbols = []
         for symbol in symbols:
-            if symbol and '-' not in symbol:
+            if symbol and "-" not in symbol:
                 symbol = f"{symbol}-USD"  # Default to USD pair
             if symbol:
                 formatted_symbols.append(symbol)
-        
+
         return [CryptoPair(symbol=symbol) for symbol in formatted_symbols]
 
     def read(self, partition: CryptoPair) -> Generator[Row, None, None]:
         """Read crypto data for a single trading pair partition."""
         symbol = partition.symbol
-        
+
         try:
             yield from self._read_crypto_pair_data(symbol)
         except Exception as e:
@@ -154,34 +156,36 @@ def _read_crypto_pair_data(self, symbol: str) -> Generator[Row, None, None]:
             # Get best bid/ask data for the trading pair using query parameters
             path = f"/api/v1/crypto/marketdata/best_bid_ask/?symbol={symbol}"
             market_data = self._make_authenticated_request("GET", path)
-            
-            if market_data and 'results' in market_data:
-                for quote in market_data['results']:
+
+            if market_data and "results" in market_data:
+                for quote in market_data["results"]:
                     # Parse numeric values safely
-                    def safe_float(value: Union[str, int, float, None], default: float = 0.0) -> float:
+                    def safe_float(
+                        value: Union[str, int, float, None], default: float = 0.0
+                    ) -> float:
                         if value is None or value == "":
                             return default
                         try:
                             return float(value)
                         except (ValueError, TypeError):
                             return default
-                    
+
                     # Extract market data fields from best bid/ask response
                     # Use the correct field names from the API response
-                    price = safe_float(quote.get('price'))
-                    bid_price = safe_float(quote.get('bid_inclusive_of_sell_spread'))
-                    ask_price = safe_float(quote.get('ask_inclusive_of_buy_spread'))
-                    
+                    price = safe_float(quote.get("price"))
+                    bid_price = safe_float(quote.get("bid_inclusive_of_sell_spread"))
+                    ask_price = safe_float(quote.get("ask_inclusive_of_buy_spread"))
+
                     yield Row(
                         symbol=symbol,
                         price=price,
                         bid_price=bid_price,
                         ask_price=ask_price,
-                        updated_at=quote.get('timestamp', "")
+                        updated_at=quote.get("timestamp", ""),
                     )
             else:
                 print(f"Warning: No market data found for {symbol}")
-                
+
         except requests.exceptions.RequestException as e:
             print(f"Network error fetching data for {symbol}: {str(e)}")
         except (ValueError, KeyError) as e:
@@ -256,10 +260,7 @@ def name(cls) -> str:
         return "robinhood"
 
     def schema(self) -> str:
-        return (
-            "symbol string, price double, bid_price double, ask_price double, "
-            "updated_at string"
-        )
+        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)
+        return RobinhoodDataReader(schema, self.options)