Refactor and type-annotate real estate ETL scripts

- Add type hints to database.py, scan_properties.py, scraper.py, and telegram_api.py
- Improve error handling for environment variable checks
- Restructure scan_properties.py with a main() function
- Update method signatures to include type annotations
- Enhance code readability and type safety

Author: Adriano Sanges

real-estate-etl/database.py

@@ -1,7 +1,8 @@
 import logging
+import polars as pl
+import duckdb
 
-
-def clean_properties(con) -> None:
+def clean_properties(con: duckdb.DuckDBPyConnection) -> None:
     logging.debug("Starting property cleaning")
     create_table_query = """
 CREATE TABLE IF NOT EXISTS main.cleaned_properties (
@@ -55,6 +56,6 @@ def clean_properties(con) -> None:
     con.sql("DELETE FROM main.properties;")
 
 
-def get_new_properties(con) -> None:
+def get_new_properties(con: duckdb.DuckDBPyConnection) -> pl.DataFrame:
     df = con.sql("SELECT * FROM main.new_properties;").pl()
     return df

real-estate-etl/scan_properties.py

@@ -1,44 +1,46 @@
+from typing import List, Dict, Any
 from bs4 import BeautifulSoup
-import polars
+import polars as pl
 import duckdb
 from scraper import parse_listing
 from database import clean_properties, get_new_properties
 from dotenv import load_dotenv
 import os
 from telegram_api import send_message, format_property_message
 
-
-
-if __name__ == "__main__":
-
-
+def main() -> None:
+    """Main function to scrape properties and send messages."""
     load_dotenv()
 
-    url = os.getenv("scrape_url")
-    warehouse_name = os.getenv("warehouse_name")
-    motherduck_token = os.getenv("motherduck_token")
+    url: str = os.getenv("scrape_url")
+    warehouse_name: str = os.getenv("warehouse_name")
+    motherduck_token: str = os.getenv("motherduck_token")
 
-    data = parse_listing(url)
+    if not url or not warehouse_name or not motherduck_token:
+        raise ValueError("Environment variables for URL, warehouse name, or token are not set.")
 
-    polars_df = polars.DataFrame(data)
+    data: List[Dict[str, Any]] = parse_listing(url)
 
-    con = duckdb.connect(f"md:{warehouse_name}?motherduck_token={motherduck_token}")
+    polars_df: pl.DataFrame = pl.DataFrame(data)
+
+    con: duckdb.DuckDBPyConnection = duckdb.connect(f"md:{warehouse_name}?motherduck_token={motherduck_token}")
 
     con.sql("create table if not exists main.properties as select * from polars_df")
 
     clean_properties(con)
 
-    new_properties = get_new_properties(con)
-    # Iterate over the DataFrame and format each property
-    messages = [format_property_message(row) for row in new_properties.iter_rows(named=True)]
+    new_properties: pl.DataFrame = get_new_properties(con)
+
+    # Format and send messages
+    messages: List[str] = [format_property_message(row) for row in new_properties.to_dicts()]
 
-# Send messages in chunks of two
+    # Send messages in chunks of two
     for i in range(0, len(messages), 2):
-        # Get the current chunk of two messages
-        message_chunk = messages[i:i+2]
-        # Join the two messages with a separator
-        full_message = "\n\n".join(message_chunk)
-        # Send the combined message
+        message_chunk: List[str] = messages[i:i+2]
+        full_message: str = "\n\n".join(message_chunk)
         send_message(full_message)
 
+if __name__ == "__main__":
+    main()
+
 

real-estate-etl/scraper.py

@@ -2,15 +2,16 @@
 import logging
 import requests
 from bs4 import BeautifulSoup
+from typing import Optional, Dict, List
 
 
-def parse_price(price_raw):
+def parse_price(price_raw: Optional[str]) -> Optional[int]:
     if not price_raw:
         return None
     price_cleaned = re.sub(r'[^\d]', '', price_raw)
     return int(price_cleaned) if price_cleaned else None
 
-def parse_page(url):
+def parse_page(url: str) -> Dict[str, Optional[any]]:
     logging.debug("Parsing page: %s", url)
     response = requests.get(url)
     soup = BeautifulSoup(response.text, 'html.parser')
@@ -30,7 +31,7 @@ def parse_page(url):
         floor_match = re.search(r'Piano\s(\d+)', soup.text)
         floor = int(floor_match.group(1)) if floor_match else None
 
-                # Find the feature item related to parking/garage
+        # Find the feature item related to parking/garage
         garage_feature = listing.find('dt', class_='re-featuresItem__title', string="Box, posti auto")
 
         if garage_feature:
@@ -55,7 +56,7 @@ def parse_page(url):
 
     return data
 
-def parse_listing(url):
+def parse_listing(url: str) -> List[Dict[str, Optional[any]]]:
     logging.debug("Fetching main listing page: %s", url)
     response = requests.get(url)
     soup = BeautifulSoup(response.text, 'html.parser')

real-estate-etl/telegram_api.py

@@ -1,19 +1,25 @@
 import telegram
 import os
-
+from typing import Dict, Any
 from dotenv import load_dotenv
 
 load_dotenv()
 
-telegram_bot_api_key = os.getenv('telegram_bot_api_key')
+TELEGRAM_BOT_API_KEY: str = os.getenv("telegram_bot_api_key")
+
+if TELEGRAM_BOT_API_KEY is None:
+    raise ValueError("TELEGRAM_BOT_API_KEY is not set. Please check your environment variables.")
+
+bot: telegram.Bot = telegram.Bot(TELEGRAM_BOT_API_KEY)
+
 chat_id = os.getenv('chat_id')
 chat_tag = os.getenv('chat_tag')
 
-bot = telegram.Bot(telegram_bot_api_key)
 
 
 # Function to format the message
-def format_property_message(row):
+def format_property_message(row: Dict[str, Any]) -> str:
+    """Format a property message for sending via Telegram."""
     return (
         f"🏠 **{row['title']}**\n"
         f"📍 Location: {row['city']}, {row['neighbourhood']}, {row['road']}\n"
@@ -25,7 +31,9 @@ def format_property_message(row):
     )
 
 
-def send_message(message):
-
+def send_message(message: str) -> None:
+    """Send a message using the Telegram bot."""
+    if chat_id is None:
+        raise ValueError("TELEGRAM_CHAT_ID is not set. Please check your environment variables.")
     bot.sendMessage(chat_id=chat_id, text=message, parse_mode=telegram.ParseMode.MARKDOWN)