Merge pull request #30 from jeafreezy/cli-bug-fix

Feat: CLI support + Enhancements

Author: denironyx

README.md

@@ -107,24 +107,24 @@ python -m graphfaker.cli --help
 
 #### Generate a Synthetic Social Graph
 ```sh
-python -m graphfaker.cli gen \
-    --source faker \
+python -m graphfaker.cli  \
+    --fetcher faker \
     --total-nodes 100 \
     --total-edges 500
 ```
 
 #### Generate a Real-World Road Network (OSM)
 ```sh
-python -m graphfaker.cli gen \
-    --source osm \
+python -m graphfaker.cli  \
+    --fetcher osm \
     --place "Berlin, Germany" \
     --network-type drive
 ```
 
 #### Generate a Flight Network (Airlines/Airports/Flights)
 ```sh
-python -m graphfaker.cli gen \
-    --source flights \
+python -m graphfaker.cli \
+    --fetcher flights \
     --country "United States" \
     --year 2024 \
     --month 1

graphfaker/cli.py

@@ -1,58 +1,72 @@
-# graphfaker/cli.py
 """
 Command-line interface for GraphFaker.
 """
+
 from venv import logger
 import typer
 from graphfaker.core import GraphFaker
+from graphfaker.enums import FetcherType
 from graphfaker.fetchers.osm import OSMGraphFetcher
 from graphfaker.fetchers.flights import FlightGraphFetcher
+from graphfaker.utils import parse_date_range
 
 app = typer.Typer()
 
 
-@app.command()
+@app.command(short_help="Generate a graph using GraphFaker.")
 def gen(
-    source: str = typer.Option("faker", help="Generation source faker library)"),
-    # for faker source
-    total_nodes: int = typer.Option(100, help="Total nodes for random mode"),
-    total_edges: int = typer.Option(1000, help="Total edges for random mode"),
-    # for osm source
+    fetcher: FetcherType = typer.Option(FetcherType.FAKER, help="Fetcher type to use."),
+    # for FetcherType.FAKER source
+    total_nodes: int = typer.Option(100, help="Total nodes for random mode."),
+    total_edges: int = typer.Option(1000, help="Total edges for random mode."),
+    # for FetcherType.OSM source
     place: str = typer.Option(
-        None, help="OSM place name (e.g., 'Soho Square, London, UK')"
+        None, help="OSM place name (e.g., 'Soho Square, London, UK')."
     ),
     address: str = typer.Option(
-        None, help="OSM address (e.g., '1600 Amphitheatre Parkway, Mountain View, CA')"
+        None, help="OSM address (e.g., '1600 Amphitheatre Parkway, Mountain View, CA.')"
     ),
-    bbox: str = typer.Option(None, help="OSM bounding box as 'north,south,east,west'"),
+    bbox: str = typer.Option(None, help="OSM bounding box as 'north,south,east,west.'"),
     network_type: str = typer.Option(
-        "drive", help="OSM network type: drive | walk | bike | all"
+        "drive", help="OSM network type: drive | walk | bike | all."
     ),
-    simplify: str = typer.Option(True, help="Simplify OSM graph topology"),
-    retain_all: bool = typer.Option(False, help="Retain all components in OSM graph"),
+    simplify: bool = typer.Option(True, help="Simplify OSM graph topology."),
+    retain_all: bool = typer.Option(False, help="Retain all components in OSM graph."),
     dist: int = typer.Option(
-        1000, help="Search radius (meters) when fetching around address"
+        1000, help="Search radius (meters) when fetching around address."
     ),
-    # for flight source
+    # for FetcherType.FLIGHT source
     country: str = typer.Option(
-        "United States", help="Filter airports by country for flight data"
+        "United States",
+        help="Filter airports by country for flight data. e.g 'United States'.",
+    ),
+    year: int = typer.Option(
+        2024, help="Year (YYYY) for single-month flight fetch. e.g. 2024."
+    ),
+    month: int = typer.Option(
+        1, help="Month (1-12) for single-month flight fetch. e.g. 1 for January."
+    ),
+    date_range: str = typer.Option(
+        None,
+        help="Year, Month and day range (YYYY-MM-DD,YYYY-MM-DD) for flight data. e.g. '2024-01-01,2024-01-15'.",
     ),
-    year: int = typer.Option(2024, help="Year (YYYY) for single-month flight fetch"),
-    month: int = typer.Option(1, help="Month (1-12) for single-month flight fetch"),
-    date_range: tuple = typer.Option(None, help="Year and Month range for flight data"),
 ):
     """Generate a graph using GraphFaker."""
     gf = GraphFaker()
-    if source == "faker":
-        G = gf.generate_graph(total_nodes=total_nodes, total_edges=total_edges)
 
-    elif source == "osm":
+    if fetcher == FetcherType.FAKER:
+
+        g = gf.generate_graph(total_nodes=total_nodes, total_edges=total_edges)
+        print(g)
+        return g
+
+    elif fetcher == FetcherType.OSM:
         # parse bbox string if provided
         bbox_tuple = None
         if bbox:
             north, south, east, west = map(float, bbox.split(","))
             bbox_tuple = (north, south, east, west)
-        G = OSMGraphFetcher.fetch_network(
+        g = OSMGraphFetcher.fetch_network(
             place=place,
             address=address,
             bbox=bbox_tuple,
@@ -61,28 +75,34 @@ def gen(
             retain_all=retain_all,
             dist=dist,
         )
-    elif source == "flights":
+        print(g)
+        return g
+    else:
+        # Flight fetcher
+        parsed_date_range = parse_date_range(date_range) if date_range else None
+
+        # validate year and month
+        if not (1 <= month <= 12):
+            raise ValueError("Month must be between 1 and 12.")
+        if not (1900 <= year <= 2100):
+            raise ValueError("Year must be between 1900 and 2100.")
 
         airlines_df = FlightGraphFetcher.fetch_airlines()
 
         airports_df = FlightGraphFetcher.fetch_airports(country=country)
 
-        # 2) Fetch on-time performance data
         flights_df = FlightGraphFetcher.fetch_flights(
-            year=year, month=month, date_range=date_range
+            year=year, month=month, date_range=parsed_date_range
         )
         logger.info(
             f"Fetched {len(airlines_df)} airlines, "
             f"{len(airports_df)} airports, "
             f"{len(flights_df)} flights."
         )
 
-        # 3) Build the NetworkX graph
-        G = FlightGraphFetcher.build_graph(airlines_df, airports_df, flights_df)
-
-    else:
-        typer.echo(f"Source '{source}' not supported.")
-        raise typer.Exit(code=1)
+        g = FlightGraphFetcher.build_graph(airlines_df, airports_df, flights_df)
+        print(g)
+        return g
 
 
 if __name__ == "__main__":

graphfaker/core.py

@@ -2,6 +2,7 @@
 A multi-domain network connecting entities across social, geographical, and commercial dimensions.
 """
 
+from typing import Optional
 import networkx as nx
 import random
 from faker import Faker
@@ -203,9 +204,9 @@ def generate_edges(self, total_edges=1000):
 
     def _generate_osm(
         self,
-        place: str = None,
-        address: str = None,
-        bbox: tuple = None,
+        place: Optional[str] = None,
+        address: Optional[str] = None,
+        bbox: Optional[tuple] = None,
         network_type: str = "drive",
         simplify: bool = True,
         retain_all: bool = False,
@@ -219,17 +220,17 @@ def _generate_osm(
             network_type=network_type,
             simplify=simplify,
             retain_all=retain_all,
-            dist=dist
+            dist=dist,
         )
         self.G = G
         return G
 
     def _generate_flights(
         self,
         country: str = "United States",
-        year: int = None,
-        month: int = None,
-        date_range: tuple = None,
+        year: Optional[int] = None,
+        month: Optional[int] = None,
+        date_range: Optional[tuple] = None,
     ):
         """
         Fetch flights, airport, and airline via FlightFetcher
@@ -250,14 +251,15 @@ def _generate_flights(
 
         G = FlightGraphFetcher.build_graph(airlines_df, airports_df, flights_df)
         self.G = G
-        return G
+
         # Inform users of which span was downloaded
         if date_range:
             start, end = date_range
             logger.info(f"Flight data covers {start} -> {end}")
 
         else:
             logger.info(f"Flight data for {year}-{month:02d}")
+        return G
 
     def _generate_faker(self, total_nodes=100, total_edges=1000):
         """Generates the complete Social Knowledge Graph."""
@@ -270,28 +272,34 @@ def generate_graph(
         source: str = "faker",
         total_nodes: int = 100,
         total_edges: int = 1000,
-        place: str = None,
-        address: str = None,
-        bbox: tuple = None,
+        place: Optional[str] = None,
+        address: Optional[str] = None,
+        bbox: Optional[tuple] = None,
         network_type: str = "drive",
         simplify: bool = True,
         retain_all: bool = False,
         dist: float = 1000,
         country: str = "United States",
         year: int = 2024,
         month: int = 1,
-        date_range: tuple = None,
+        date_range: Optional[tuple] = None,
     ) -> nx.DiGraph:
         """
         Unified entrypoint: choose 'random' or 'osm'.
         Pass kwargs depending on source.
         """
+
         if source == "faker":
             return self._generate_faker(
-                total_nodes=total_nodes,
-                total_edges=total_edges
+                total_nodes=total_nodes, total_edges=total_edges
             )
         elif source == "osm":
+            logger.info(
+                f"Generating OSM graph with source={source}, "
+                f"place={place}, address={address}, bbox={bbox}, "
+                f"network_type={network_type}, simplify={simplify}, "
+                f"retain_all={retain_all}, dist={dist}"
+            )
             return self._generate_osm(
                 place=place,
                 address=address,

graphfaker/enums.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class FetcherType(str, Enum):
+    """Enum for different fetcher types."""
+
+    OSM = "osm"
+    FLIGHTS = "flights"
+    FAKER = "faker"

graphfaker/fetchers/flights.py

@@ -33,7 +33,7 @@
 import pandas as pd
 from tqdm.auto import tqdm
 import networkx as nx
-import time
+
 
 # suppress only the single warning from unverified HTTPS
 import urllib3
@@ -124,6 +124,7 @@ def fetch_airlines() -> pd.DataFrame:
         Raises:
             HTTPError if download fails.
         """
+        logger.info("Fetching airlines lookup from BTS…")
         resp = requests.get(AIRLINE_LOOKUP_URL, verify=False)
         resp.raise_for_status()
         df = pd.read_csv(StringIO(resp.text))
@@ -145,6 +146,7 @@ def fetch_airports(
         Returns:
             pd.DataFrame with columns ['faa','name','city','country','lat','lon']
         """
+        logger.info("Fetching airports dataset from OpenFlights…")
         df = pd.read_csv(
             AIRPORTS_URL,
             header=None,
@@ -220,6 +222,10 @@ def fetch_flights(
         Raises:
             ValueError if neither valid year/month nor date_range provided.
         """
+        logger.info(
+            f"Fetching flight performance data for {year}-{month:02d} "
+            f"or date range {date_range}…"
+        )
 
         def load_month(y, m):
             url = f"https://transtats.bts.gov/PREZIP/On_Time_Reporting_Carrier_On_Time_Performance_1987_present_{y}_{m}.zip"